// Package tunnel defines the framed multiplexing protocol used for DIRECT port
// forwarding over the dedicated host port (default 8446). Unlike the HTTP
// long-poll forward, this is a raw TCP stream the agent dials outbound; it
// carries many concurrent forwarded connections (streams) for one agent and
// has far lower latency than long-polling.
//
// Frame layout (big-endian):
//
//	type(1) | stream(4) | length(4) | payload(length)
package tunnel

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"io"
)

// DefaultPort is the host port for the direct tunnel (TCP and UDP).
const DefaultPort = 8446

// Frame types.
const (
	TypeOpen  byte = 1 // host->agent: open a TCP stream to 127.0.0.1:port (payload = port uint16)
	TypeData  byte = 2 // both: stream bytes
	TypeClose byte = 3 // both: close the stream
	TypePing  byte = 4 // keepalive
)

// MaxPayload caps a single frame payload (1 MiB).
const MaxPayload = 1 << 20

// Frame is one multiplexed message.
type Frame struct {
	Type    byte
	Stream  uint32
	Payload []byte
}

// WriteFrame writes a frame. The caller must serialize writes to one conn.
func WriteFrame(w io.Writer, f Frame) error {
	var hdr [9]byte
	hdr[0] = f.Type
	binary.BigEndian.PutUint32(hdr[1:5], f.Stream)
	binary.BigEndian.PutUint32(hdr[5:9], uint32(len(f.Payload)))
	if _, err := w.Write(hdr[:]); err != nil {
		return err
	}
	if len(f.Payload) > 0 {
		_, err := w.Write(f.Payload)
		return err
	}
	return nil
}

// ReadFrame reads one frame.
func ReadFrame(r io.Reader) (Frame, error) {
	var hdr [9]byte
	if _, err := io.ReadFull(r, hdr[:]); err != nil {
		return Frame{}, err
	}
	n := binary.BigEndian.Uint32(hdr[5:9])
	if n > MaxPayload {
		return Frame{}, io.ErrUnexpectedEOF
	}
	f := Frame{Type: hdr[0], Stream: binary.BigEndian.Uint32(hdr[1:5])}
	if n > 0 {
		f.Payload = make([]byte, n)
		if _, err := io.ReadFull(r, f.Payload); err != nil {
			return Frame{}, err
		}
	}
	return f, nil
}

// AuthKey derives the shared HMAC key from an agent token: sha256(token).
// The server holds the same value as the stored token hash, so neither side
// ever puts the plaintext token on the wire.
func AuthKey(token string) []byte {
	sum := sha256.Sum256([]byte(token))
	return sum[:]
}

// KeyFromHash decodes a stored hex token hash into the HMAC key bytes.
func KeyFromHash(tokenHashHex string) ([]byte, error) {
	return hex.DecodeString(tokenHashHex)
}

// Sign returns hex(HMAC-SHA256(key, nonce)).
func Sign(key, nonce []byte) string {
	m := hmac.New(sha256.New, key)
	m.Write(nonce)
	return hex.EncodeToString(m.Sum(nil))
}

// EqualSig compares two hex signatures in constant time.
func EqualSig(a, b string) bool {
	return hmac.Equal([]byte(a), []byte(b))
}