mirror of
https://github.com/k3s-io/k3s.git
synced 2024-06-07 19:41:36 +00:00
99 lines
2.4 KiB
Go
99 lines
2.4 KiB
Go
package remotedialer
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
"github.com/pkg/errors"
|
|
"github.com/sirupsen/logrus"
|
|
)
|
|
|
|
var (
|
|
errFailedAuth = errors.New("failed authentication")
|
|
errWrongMessageType = errors.New("wrong websocket message type")
|
|
)
|
|
|
|
type Authorizer func(req *http.Request) (clientKey string, authed bool, err error)
|
|
type ErrorWriter func(rw http.ResponseWriter, req *http.Request, code int, err error)
|
|
|
|
func DefaultErrorWriter(rw http.ResponseWriter, req *http.Request, code int, err error) {
|
|
rw.Write([]byte(err.Error()))
|
|
rw.WriteHeader(code)
|
|
}
|
|
|
|
type Server struct {
|
|
PeerID string
|
|
PeerToken string
|
|
authorizer Authorizer
|
|
errorWriter ErrorWriter
|
|
sessions *sessionManager
|
|
peers map[string]peer
|
|
peerLock sync.Mutex
|
|
}
|
|
|
|
func New(auth Authorizer, errorWriter ErrorWriter) *Server {
|
|
return &Server{
|
|
peers: map[string]peer{},
|
|
authorizer: auth,
|
|
errorWriter: errorWriter,
|
|
sessions: newSessionManager(),
|
|
}
|
|
}
|
|
|
|
func (s *Server) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
|
|
clientKey, authed, peer, err := s.auth(req)
|
|
if err != nil {
|
|
s.errorWriter(rw, req, 400, err)
|
|
return
|
|
}
|
|
if !authed {
|
|
s.errorWriter(rw, req, 401, errFailedAuth)
|
|
return
|
|
}
|
|
|
|
logrus.Infof("Handling backend connection request [%s]", clientKey)
|
|
|
|
upgrader := websocket.Upgrader{
|
|
HandshakeTimeout: 5 * time.Second,
|
|
CheckOrigin: func(r *http.Request) bool { return true },
|
|
Error: s.errorWriter,
|
|
}
|
|
|
|
wsConn, err := upgrader.Upgrade(rw, req, nil)
|
|
if err != nil {
|
|
s.errorWriter(rw, req, 400, errors.Wrapf(err, "Error during upgrade for host [%v]", clientKey))
|
|
return
|
|
}
|
|
|
|
session := s.sessions.add(clientKey, wsConn, peer)
|
|
defer s.sessions.remove(session)
|
|
|
|
// Don't need to associate req.Context() to the Session, it will cancel otherwise
|
|
code, err := session.Serve(context.Background())
|
|
if err != nil {
|
|
// Hijacked so we can't write to the client
|
|
logrus.Infof("error in remotedialer server [%d]: %v", code, err)
|
|
}
|
|
}
|
|
|
|
func (s *Server) auth(req *http.Request) (clientKey string, authed, peer bool, err error) {
|
|
id := req.Header.Get(ID)
|
|
token := req.Header.Get(Token)
|
|
if id != "" && token != "" {
|
|
// peer authentication
|
|
s.peerLock.Lock()
|
|
p, ok := s.peers[id]
|
|
s.peerLock.Unlock()
|
|
|
|
if ok && p.token == token {
|
|
return id, true, true, nil
|
|
}
|
|
}
|
|
|
|
id, authed, err = s.authorizer(req)
|
|
return id, authed, false, err
|
|
}
|