Feat: ConnAuth HMAC verification gated by YAMA_CONN_AUTH_KEY

Verify the HMAC-SHA256 signature and timestamp of device sub-connections
(screen/terminal/file) when the YAMA_CONN_AUTH_KEY environment variable is
set. With the key set, a sub-connection whose signature or clock skew fails
is answered with a non-OK ConnAuth status and dropped; without the key the
legacy always-OK behavior is preserved.

The key is never compiled in: it is read from the environment at startup, so
no secret enters the repository. Source-IP binding in hub.BindScreenConn
remains in place as defense-in-depth.

Co-Authored-By: deepseek-v4-pro

server/go/cmd/main.go

@@ -2,7 +2,10 @@ package main
 
 import (
 	"context"
+	"crypto/hmac"
+	"crypto/sha256"
 	"encoding/binary"
+	"encoding/hex"
 	"flag"
 	"fmt"
 	"net/http"
@@ -26,12 +29,13 @@ import (
 
 // MyHandler implements the server.Handler interface
 type MyHandler struct {
-	log        *logger.Logger
-	auth       *auth.Authenticator
-	srv        *server.Server
-	hub        *hub.Hub
-	signer     licensing.Signer // CMD_MASTERSETTING signer (local / remote / noop)
-	loginLimit *wsauth.RateLimiter // per-IP TOKEN_AUTH throttle
+	log         *logger.Logger
+	auth        *auth.Authenticator
+	srv         *server.Server
+	hub         *hub.Hub
+	signer      licensing.Signer    // CMD_MASTERSETTING signer (local / remote / noop)
+	loginLimit  *wsauth.RateLimiter // per-IP TOKEN_AUTH throttle
+	connAuthKey []byte              // nil = ConnAuth HMAC verification disabled (legacy always-OK)
 }
 
 // OnConnect is called when a client connects
@@ -154,12 +158,25 @@ func (h *MyHandler) OnReceive(ctx *connection.Context, data []byte) {
 // on our 256-byte ConnAuthAck. Without an OK reply the client closes the
 // connection, so a missing ack here means nothing else can proceed.
 //
-// The handshake includes an HMAC signature field. The reference server
-// treats verification failures as soft (logs and still allows commands),
-// and the signing primitive lives in a vendored component out of scope
-// for this server, so we always reply OK and let TOKEN_BITMAPINFO carry
-// the device ID via offset 41 when the screen sub-conn proceeds.
+// When the operator supplied YAMA_CONN_AUTH_KEY, the packet's HMAC-SHA256
+// signature and timestamp are verified first (see verifyConnAuth); a bad
+// packet is answered with a non-OK status and the sub-conn is dropped.
+// Without a key we keep the legacy behavior: always reply OK and let
+// TOKEN_BITMAPINFO carry the device ID via offset 41 when the screen
+// sub-conn proceeds.
 func (h *MyHandler) handleConnAuth(ctx *connection.Context, data []byte) {
+	// When the operator supplied a ConnAuth HMAC key (YAMA_CONN_AUTH_KEY),
+	// verify the packet's signature and timestamp before trusting its
+	// clientID. This closes the sub-connection hijack hole that the
+	// source-IP check in hub.BindScreenConn otherwise alone guards against.
+	// A nil key preserves the legacy always-OK behavior.
+	if len(h.connAuthKey) > 0 {
+		if status := h.verifyConnAuth(data); status != protocol.ConnAuthStatusOK {
+			h.rejectConnAuth(ctx, status)
+			return
+		}
+	}
+
 	// Pin the parent device's ClientID onto the sub-conn. Without this,
 	// later 1-byte tokens (TOKEN_TERMINAL_START / TOKEN_SHELL_START) have
 	// no way to identify which device they belong to — they carry no
@@ -188,6 +205,61 @@ func (h *MyHandler) handleConnAuth(ctx *connection.Context, data []byte) {
 	}
 }
 
+// verifyConnAuth checks the ConnAuthPacket's HMAC-SHA256 signature and its
+// timestamp against the configured shared key. Returns a ConnAuthStatus code
+// (OK on success).
+func (h *MyHandler) verifyConnAuth(data []byte) byte {
+	// Need the full [token][clientID][timestamp][nonce][signature] prefix.
+	if len(data) < protocol.ConnAuthOffSignature+protocol.ConnAuthSignatureSize {
+		return protocol.ConnAuthStatusBadSize
+	}
+
+	// Anti-replay: the client signs its own Unix timestamp, so an old (but
+	// valid) packet cannot be replayed once it leaves the tolerance window.
+	ts := binary.LittleEndian.Uint64(
+		data[protocol.ConnAuthOffTimestamp : protocol.ConnAuthOffTimestamp+8])
+	now := uint64(time.Now().Unix())
+	var skew uint64
+	if now >= ts {
+		skew = now - ts
+	} else {
+		skew = ts - now
+	}
+	if skew > protocol.ConnAuthTimestampToleranceSec {
+		return protocol.ConnAuthStatusClockSkew
+	}
+
+	// HMAC-SHA256 over clientID||timestamp||nonce (offsets 1..33), matching
+	// the client's signMessage("", sigInput, 32). The client stores the hex
+	// encoding of the 32-byte digest in the 64-byte signature field.
+	mac := hmac.New(sha256.New, h.connAuthKey)
+	mac.Write(data[protocol.ConnAuthOffClientID:protocol.ConnAuthOffSignature])
+	want := make([]byte, protocol.ConnAuthSignatureSize)
+	hex.Encode(want, mac.Sum(nil))
+	if !hmac.Equal(want, data[protocol.ConnAuthOffSignature:protocol.ConnAuthOffSignature+protocol.ConnAuthSignatureSize]) {
+		return protocol.ConnAuthStatusBadSignature
+	}
+	return protocol.ConnAuthStatusOK
+}
+
+// rejectConnAuth replies with a non-OK ConnAuth status and closes the
+// sub-connection. The client (IOCPClient::TryHandleAuthResponse) treats any
+// non-zero status as failure and disconnects, so the close here is
+// belt-and-suspenders for peers that don't.
+func (h *MyHandler) rejectConnAuth(ctx *connection.Context, status byte) {
+	ack := make([]byte, protocol.ConnAuthAckSize)
+	ack[0] = protocol.TokenConnAuth
+	ack[protocol.ConnAuthAckOffStatus] = status
+	binary.LittleEndian.PutUint64(
+		ack[protocol.ConnAuthAckOffServerTime:protocol.ConnAuthAckOffServerTime+8],
+		uint64(time.Now().Unix()))
+	if err := h.srv.Send(ctx, ack); err != nil {
+		h.log.Error("ConnAuth reject send failed for conn=%d: %v", ctx.ID, err)
+	}
+	h.log.Warn("ConnAuth rejected conn=%d ip=%s status=%d", ctx.ID, ctx.GetPeerIP(), status)
+	_ = ctx.Close()
+}
+
 // deviceIDOfSubConn resolves the parent device of a sub-conn from the
 // ClientID pinned by handleConnAuth. Returns "" for the rare case of a
 // legacy client that skipped ConnAuth (the Go server's only target is
@@ -933,6 +1005,18 @@ func main() {
 		log.Info("TLS disabled (YAMA_TLS_CERT/YAMA_TLS_KEY not set) — device port is plaintext")
 	}
 
+	// Optional ConnAuth HMAC verification for sub-connections (screen/terminal/
+	// file). The shared key must match what the client signs with — either the
+	// client's built-in HMAC key (see client license.cpp / license_unix.cpp) or
+	// a custom key once the client is taught to read one. Unset means legacy
+	// behavior: accept any sub-connection and rely on the source-IP check in
+	// hub.BindScreenConn alone.
+	var connAuthKey []byte
+	if v := strings.TrimSpace(os.Getenv("YAMA_CONN_AUTH_KEY")); v != "" {
+		connAuthKey = []byte(v)
+		log.Info("ConnAuth HMAC verification enabled for device sub-connections (YAMA_CONN_AUTH_KEY)")
+	}
+
 	// Create servers for each port
 	var servers []*server.Server
 	for _, port := range ports {
@@ -949,12 +1033,13 @@ func main() {
 
 		// Create handler for this server
 		handler := &MyHandler{
-			log:        log.WithPrefix(fmt.Sprintf("Handler:%d", port)),
-			auth:       authenticator,
-			srv:        srv,
-			hub:        deviceHub,
-			signer:     signer,
-			loginLimit: authLoginLimit,
+			log:         log.WithPrefix(fmt.Sprintf("Handler:%d", port)),
+			auth:        authenticator,
+			srv:         srv,
+			hub:         deviceHub,
+			signer:      signer,
+			loginLimit:  authLoginLimit,
+			connAuthKey: connAuthKey,
 		}
 		srv.SetHandler(handler)