Files
redapricot/client/wire/framedconn.go
T
2026-07-25 16:33:28 +08:00

108 lines
3.0 KiB
Go

package wire
import (
"bufio"
"errors"
"io"
"net"
"sync"
"time"
"golang.org/x/crypto/chacha20"
)
// MaxFrame is the maximum decrypted frame payload size (1 MiB).
const MaxFrame = 1 << 20
// WriteTimeout bounds a single frame write. A peer that stops reading must not
// be able to park every stream on the connection inside WriteFrame forever: the
// write mutex is held for the whole socket write, so one stalled write would
// otherwise wedge the entire multiplexed connection.
const WriteTimeout = 30 * time.Second
var (
errFrameTooBig = errors.New("wire: frame exceeds max size")
// ErrBroken is returned once a write has failed. The ChaCha20 keystream has
// already advanced (and the socket may hold a partial frame), so the
// connection can never be resynchronized and is closed for good.
ErrBroken = errors.New("wire: connection is broken")
)
// FramedConn is the encrypted, length-prefixed frame transport (PROTOCOL.md §3.1).
// The VarInt length prefix is plaintext; the payload is ChaCha20-encrypted with a
// continuous per-direction keystream. Writes are serialized; reads are expected
// from a single goroutine.
type FramedConn struct {
conn net.Conn
r *bufio.Reader
in *chacha20.Cipher
out *chacha20.Cipher
wmu sync.Mutex
broken bool
}
func NewFramedConn(conn net.Conn, in, out *chacha20.Cipher) *FramedConn {
return &FramedConn{
conn: conn,
r: bufio.NewReader(conn),
in: in,
out: out,
}
}
// SwitchCiphers swaps both ciphers at a frame boundary (Phase A → Phase B).
// Only call this from the same goroutine sequence as reads/writes during the
// handshake, before concurrency begins.
func (f *FramedConn) SwitchCiphers(in, out *chacha20.Cipher) {
f.in = in
f.out = out
}
// ReadFrame reads and decrypts one frame payload.
func (f *FramedConn) ReadFrame() ([]byte, error) {
n, err := ReadVarInt(f.r)
if err != nil {
return nil, err
}
if n < 0 || n > MaxFrame {
return nil, errFrameTooBig
}
ct := make([]byte, n)
if _, err := io.ReadFull(f.r, ct); err != nil {
return nil, err
}
f.in.XORKeyStream(ct, ct) // decrypt in place
return ct, nil
}
// WriteFrame encrypts and sends one frame payload. Safe for concurrent callers.
//
// The write is bounded by WriteTimeout. On any write error the connection is
// marked broken and closed, which unblocks the reader so the owner can tear the
// session down instead of leaving every stream parked on the write mutex.
func (f *FramedConn) WriteFrame(payload []byte) error {
f.wmu.Lock()
defer f.wmu.Unlock()
if f.broken {
return ErrBroken
}
ct := make([]byte, len(payload))
f.out.XORKeyStream(ct, payload)
out := AppendVarInt(make([]byte, 0, VarIntMaxBytes+len(ct)), len(ct))
out = append(out, ct...)
_ = f.conn.SetWriteDeadline(time.Now().Add(WriteTimeout))
_, err := f.conn.Write(out)
if err != nil {
f.broken = true
_ = f.conn.Close()
return err
}
_ = f.conn.SetWriteDeadline(time.Time{})
return nil
}
func (f *FramedConn) Close() error { return f.conn.Close() }
func (f *FramedConn) RemoteAddr() net.Addr { return f.conn.RemoteAddr() }