diff --git a/muxterm/.gitignore b/muxterm/.gitignore new file mode 100644 index 0000000..4a1a86d --- /dev/null +++ b/muxterm/.gitignore @@ -0,0 +1 @@ +muxterm \ No newline at end of file diff --git a/muxterm/go.mod b/muxterm/go.mod index 09cf5f4..e9d0c5f 100644 --- a/muxterm/go.mod +++ b/muxterm/go.mod @@ -3,3 +3,5 @@ module github.com/muxplex/muxterm go 1.22 require github.com/creack/pty v1.1.24 + +require github.com/gorilla/websocket v1.5.3 // indirect diff --git a/muxterm/go.sum b/muxterm/go.sum index 5310a43..a15fd45 100644 --- a/muxterm/go.sum +++ b/muxterm/go.sum @@ -1,2 +1,4 @@ github.com/creack/pty v1.1.24 h1:bJrF4RRfyJnbTJqzRLHzcGaZK1NeM5kTC9jGgovnR1s= github.com/creack/pty v1.1.24/go.mod h1:08sCNb52WyoAwi2QDyzUCTgcvVFhUzewun7wtTfvcwE= +github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= +github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= diff --git a/muxterm/ws.go b/muxterm/ws.go new file mode 100644 index 0000000..b2ddf1f --- /dev/null +++ b/muxterm/ws.go @@ -0,0 +1,188 @@ +package main + +import ( + "crypto/hmac" + "crypto/sha256" + "encoding/hex" + "io" + "log" + "net/http" + "os" + "strconv" + "strings" + "sync" + "syscall" + "time" + + "github.com/gorilla/websocket" +) + +// ValidateToken checks an HMAC-SHA256 signed token of the form "signature.timestamp". +// It validates the timestamp is within ttl seconds (not too old) and within 5 seconds +// of clock skew (not too far in the future), then verifies the HMAC signature. +func ValidateToken(token, secret string, ttl int64) bool { + parts := strings.SplitN(token, ".", 2) + if len(parts) != 2 { + return false + } + sig := parts[0] + tsStr := parts[1] + + ts, err := strconv.ParseInt(tsStr, 10, 64) + if err != nil { + return false + } + + now := time.Now().Unix() + if now-ts > ttl { + return false + } + if ts-now > 5 { + return false + } + + mac := hmac.New(sha256.New, []byte(secret)) + mac.Write([]byte(tsStr)) + expected := hex.EncodeToString(mac.Sum(nil)) + + return hmac.Equal([]byte(sig), []byte(expected)) +} + +var upgrader = websocket.Upgrader{ + CheckOrigin: func(r *http.Request) bool { + host := r.Host + if strings.Contains(host, ":") { + host = strings.Split(host, ":")[0] + } + return host == "localhost" || host == "127.0.0.1" + }, +} + +// HandleWebSocket validates the token query parameter and upgrades the connection +// to a WebSocket, then runs the main read loop for control and binary messages. +func HandleWebSocket(pool *Pool, secret string, w http.ResponseWriter, r *http.Request) { + token := r.URL.Query().Get("token") + if !ValidateToken(token, secret, 30) { + http.Error(w, "forbidden", http.StatusForbidden) + return + } + + conn, err := upgrader.Upgrade(w, r, nil) + if err != nil { + log.Printf("ws upgrade: %v", err) + return + } + defer conn.Close() + + var ( + mu sync.Mutex + activeSession string + activePTY *os.File + ptyCancel chan struct{} + ) + + startRelay := func(sess *Session, name string) { + cancel := make(chan struct{}) + mu.Lock() + ptyCancel = cancel + mu.Unlock() + + go func() { + buf := make([]byte, 32*1024) + for { + select { + case <-cancel: + return + default: + } + n, err := sess.PTY.Read(buf) + if n > 0 { + if writeErr := conn.WriteMessage(websocket.BinaryMessage, buf[:n]); writeErr != nil { + return + } + } + if err != nil { + if err == io.EOF || strings.Contains(err.Error(), "input/output error") { + conn.WriteMessage(websocket.TextMessage, MarshalExited(name)) + } + return + } + } + }() + } + + stopRelay := func() { + mu.Lock() + defer mu.Unlock() + if ptyCancel != nil { + close(ptyCancel) + ptyCancel = nil + } + } + + defer stopRelay() + + for { + msgType, data, err := conn.ReadMessage() + if err != nil { + return + } + + switch msgType { + case websocket.TextMessage: + kind, msg := ParseControlMessage(data) + switch kind { + case "attach": + attach := msg.(*AttachMsg) + stopRelay() + + sess, err := pool.Attach(attach.Attach, 80, 24) + if err != nil { + conn.WriteMessage(websocket.TextMessage, MarshalError(err.Error())) + continue + } + + mu.Lock() + activeSession = attach.Attach + activePTY = sess.PTY + mu.Unlock() + + sendSIGWINCH(sess) + startRelay(sess, attach.Attach) + conn.WriteMessage(websocket.TextMessage, MarshalAttached(attach.Attach)) + + case "resize": + resize := msg.(*ResizeMsg) + mu.Lock() + name := activeSession + mu.Unlock() + if name != "" { + pool.Resize(name, uint16(resize.Resize.Cols), uint16(resize.Resize.Rows)) + } + + case "detach": + stopRelay() + mu.Lock() + activeSession = "" + activePTY = nil + mu.Unlock() + } + + case websocket.BinaryMessage: + mu.Lock() + pty := activePTY + mu.Unlock() + if pty != nil { + pty.Write(data) + } + } + } +} + +// sendSIGWINCH sends SIGWINCH to the session's process to trigger tmux repaint. +func sendSIGWINCH(s *Session) { + if s == nil || s.Cmd == nil || s.Cmd.Process == nil { + return + } + s.Cmd.Process.Signal(syscall.SIGWINCH) +} \ No newline at end of file diff --git a/muxterm/ws_stub.go b/muxterm/ws_stub.go deleted file mode 100644 index 7df95e3..0000000 --- a/muxterm/ws_stub.go +++ /dev/null @@ -1,8 +0,0 @@ -package main - -import "net/http" - -// HandleWebSocket is a placeholder that will be replaced in a later task. -func HandleWebSocket(pool *Pool, secret string, w http.ResponseWriter, r *http.Request) { - http.Error(w, "not implemented", http.StatusNotImplemented) -} diff --git a/muxterm/ws_test.go b/muxterm/ws_test.go new file mode 100644 index 0000000..e300e31 --- /dev/null +++ b/muxterm/ws_test.go @@ -0,0 +1,171 @@ +package main + +import ( + "crypto/hmac" + "crypto/sha256" + "encoding/hex" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/gorilla/websocket" +) + +// makeToken creates a valid HMAC-SHA256 token: sig.timestamp +func makeToken(secret string, ts int64) string { + tsStr := fmt.Sprintf("%d", ts) + mac := hmac.New(sha256.New, []byte(secret)) + mac.Write([]byte(tsStr)) + sig := hex.EncodeToString(mac.Sum(nil)) + return sig + "." + tsStr +} + +func TestValidateToken_ValidToken(t *testing.T) { + secret := "test-secret" + ts := time.Now().Unix() + token := makeToken(secret, ts) + + if !ValidateToken(token, secret, 30) { + t.Error("expected valid token to pass validation") + } +} + +func TestValidateToken_ExpiredToken(t *testing.T) { + secret := "test-secret" + ts := time.Now().Unix() - 60 // 60 seconds ago, TTL is 30 + token := makeToken(secret, ts) + + if ValidateToken(token, secret, 30) { + t.Error("expected expired token to fail validation") + } +} + +func TestValidateToken_FutureToken_WithinSkew(t *testing.T) { + secret := "test-secret" + ts := time.Now().Unix() + 3 // 3 seconds in future, within 5s skew + token := makeToken(secret, ts) + + if !ValidateToken(token, secret, 30) { + t.Error("expected near-future token within clock skew to pass") + } +} + +func TestValidateToken_FutureToken_BeyondSkew(t *testing.T) { + secret := "test-secret" + ts := time.Now().Unix() + 10 // 10 seconds in future, beyond 5s skew + token := makeToken(secret, ts) + + if ValidateToken(token, secret, 30) { + t.Error("expected far-future token beyond clock skew to fail") + } +} + +func TestValidateToken_BadSignature(t *testing.T) { + ts := time.Now().Unix() + token := makeToken("right-secret", ts) + + if ValidateToken(token, "wrong-secret", 30) { + t.Error("expected wrong secret to fail validation") + } +} + +func TestValidateToken_MalformedToken(t *testing.T) { + if ValidateToken("not-a-valid-token", "secret", 30) { + t.Error("expected malformed token (no dot) to fail") + } + if ValidateToken("abc.def", "secret", 30) { + t.Error("expected non-numeric timestamp to fail") + } + if ValidateToken("", "secret", 30) { + t.Error("expected empty token to fail") + } +} + +func TestHandleWebSocket_RejectsNoToken(t *testing.T) { + pool := NewPool() + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + HandleWebSocket(pool, "test-secret", w, r) + })) + defer server.Close() + + // Try to connect without a token + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "?token=" + _, resp, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err == nil { + t.Error("expected connection to fail without valid token") + } + if resp != nil && resp.StatusCode != http.StatusForbidden { + t.Errorf("expected 403 Forbidden, got %d", resp.StatusCode) + } +} + +func TestHandleWebSocket_RejectsExpiredToken(t *testing.T) { + pool := NewPool() + secret := "test-secret" + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + HandleWebSocket(pool, secret, w, r) + })) + defer server.Close() + + ts := time.Now().Unix() - 60 + token := makeToken(secret, ts) + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "?token=" + token + _, resp, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err == nil { + t.Error("expected connection to fail with expired token") + } + if resp != nil && resp.StatusCode != http.StatusForbidden { + t.Errorf("expected 403 Forbidden, got %d", resp.StatusCode) + } +} + +func TestHandleWebSocket_AcceptsValidToken(t *testing.T) { + pool := NewPool() + secret := "test-secret" + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + HandleWebSocket(pool, secret, w, r) + })) + defer server.Close() + + ts := time.Now().Unix() + token := makeToken(secret, ts) + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "?token=" + token + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("expected connection to succeed with valid token, got: %v", err) + } + defer conn.Close() +} + +func TestHandleWebSocket_DetachMessage(t *testing.T) { + pool := NewPool() + secret := "test-secret" + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + HandleWebSocket(pool, secret, w, r) + })) + defer server.Close() + + ts := time.Now().Unix() + token := makeToken(secret, ts) + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "?token=" + token + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + if err != nil { + t.Fatalf("expected connection to succeed, got: %v", err) + } + defer conn.Close() + + // Send detach - should not crash, just clear state + err = conn.WriteMessage(websocket.TextMessage, []byte(`{"detach":true}`)) + if err != nil { + t.Fatalf("failed to send detach message: %v", err) + } +} + +func TestSendSIGWINCH_NilProcess(t *testing.T) { + // Should not panic with nil Cmd or nil Process + sess := &Session{} + sendSIGWINCH(sess) // should not panic +} \ No newline at end of file