diff --git a/client/internal/pqkem/message.go b/client/internal/pqkem/message.go new file mode 100644 index 000000000..341f6b5f8 --- /dev/null +++ b/client/internal/pqkem/message.go @@ -0,0 +1,135 @@ +package pqkem + +import ( + "crypto/mlkem" + "fmt" +) + +// Wire framing for the PQ-KEM exchange. Messages are self-contained, versioned, +// transport-agnostic byte blobs: the same bytes ride the Signal offer/answer +// (initial, pre-tunnel) or a WireGuard-tunnel packet (rekey). They are NOT a gRPC +// service — the network layer only sees opaque []byte. +// +// Layout (all messages): [type:1][version:1][exchangeID:16][payload...] + +const ( + // ProtocolVersion is bumped on any wire-incompatible change; a peer rejects + // messages it does not understand rather than misparsing them. + ProtocolVersion uint8 = 1 + + // ExchangeIDSize identifies one offer/answer/confirm round so stale answers + // (e.g. an answer to a pre-restart offer) are dropped instead of applied. + ExchangeIDSize = 16 + + headerSize = 1 + 1 + ExchangeIDSize +) + +// MsgType tags the three message kinds of the exchange. +type MsgType uint8 + +const ( + MsgOffer MsgType = iota + 1 + MsgAnswer + MsgConfirm +) + +// ExchangeID is the per-round correlator echoed by the answer and the confirm. +type ExchangeID [ExchangeIDSize]byte + +// OfferMsg carries the initiator's public material (X25519 pub ‖ ML-KEM encap key). +type OfferMsg struct { + ExchangeID ExchangeID + // KEMOffer is the raw Initiator.Offer() blob (OfferSize bytes). + KEMOffer []byte +} + +// AnswerMsg carries the responder's reply (ML-KEM ciphertext ‖ X25519 pub) for the +// round identified by ExchangeID. +type AnswerMsg struct { + ExchangeID ExchangeID + // KEMAnswer is the raw Respond() answer blob (AnswerSize bytes). + KEMAnswer []byte +} + +// ConfirmMsg is sent ONE WAY, initiator -> responder, over the tunnel under the +// NEW PSK, for the round identified by ExchangeID. It carries no key material. +// +// It is one-way because of the knowledge asymmetry of a KEM: the initiator learns +// the responder holds the key simply by receiving the answer (the responder had to +// derive it to encapsulate), so the initiator needs no confirmation. Only the +// responder is left unsure whether the initiator received the answer and committed +// the key — this message resolves that. As a bonus, being sent under the new PSK it +// triggers the WireGuard handshake with the new key: the responder converges on +// receiving it, and the initiator converges by observing that handshake succeed +// (which fails on a PSK mismatch, so success proves the responder also committed). +type ConfirmMsg struct { + ExchangeID ExchangeID +} + +// Encode serialises the offer with its framed header. +func (m *OfferMsg) Encode() ([]byte, error) { + if len(m.KEMOffer) != OfferSize { + return nil, fmt.Errorf("offer payload: got %d, want %d", len(m.KEMOffer), OfferSize) + } + return frame(MsgOffer, m.ExchangeID, m.KEMOffer), nil +} + +// Encode serialises the answer with its framed header. +func (m *AnswerMsg) Encode() ([]byte, error) { + if len(m.KEMAnswer) != AnswerSize { + return nil, fmt.Errorf("answer payload: got %d, want %d", len(m.KEMAnswer), AnswerSize) + } + return frame(MsgAnswer, m.ExchangeID, m.KEMAnswer), nil +} + +// Encode serialises the confirm with its framed header. +func (m *ConfirmMsg) Encode() []byte { + return frame(MsgConfirm, m.ExchangeID, nil) +} + +// Decode parses a framed message into one of *OfferMsg / *AnswerMsg / *ConfirmMsg. +func Decode(buf []byte) (MsgType, any, error) { + if len(buf) < headerSize { + return 0, nil, fmt.Errorf("message too short: %d bytes", len(buf)) + } + typ := MsgType(buf[0]) + if ver := buf[1]; ver != ProtocolVersion { + return typ, nil, fmt.Errorf("unsupported protocol version %d (want %d)", ver, ProtocolVersion) + } + + var id ExchangeID + copy(id[:], buf[2:headerSize]) + payload := buf[headerSize:] + + switch typ { + case MsgOffer: + if len(payload) != OfferSize { + return typ, nil, fmt.Errorf("offer payload: got %d, want %d", len(payload), OfferSize) + } + return typ, &OfferMsg{ExchangeID: id, KEMOffer: payload}, nil + case MsgAnswer: + if len(payload) != AnswerSize { + return typ, nil, fmt.Errorf("answer payload: got %d, want %d", len(payload), AnswerSize) + } + return typ, &AnswerMsg{ExchangeID: id, KEMAnswer: payload}, nil + case MsgConfirm: + if len(payload) != 0 { + return typ, nil, fmt.Errorf("confirm payload must be empty, got %d bytes", len(payload)) + } + return typ, &ConfirmMsg{ExchangeID: id}, nil + default: + return typ, nil, fmt.Errorf("unknown message type %d", typ) + } +} + +func frame(typ MsgType, id ExchangeID, payload []byte) []byte { + buf := make([]byte, headerSize+len(payload)) + buf[0] = byte(typ) + buf[1] = ProtocolVersion + copy(buf[2:], id[:]) + copy(buf[headerSize:], payload) + return buf +} + +// compile-time assurance the KEM blob sizes referenced here stay in sync with kem.go. +var _ = [1]struct{}{}[OfferSize-(32+mlkem.EncapsulationKeySize768)] diff --git a/client/internal/pqkem/message_test.go b/client/internal/pqkem/message_test.go new file mode 100644 index 000000000..a51135ea0 --- /dev/null +++ b/client/internal/pqkem/message_test.go @@ -0,0 +1,61 @@ +package pqkem + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestMessageRoundTrip(t *testing.T) { + init, err := NewInitiator() + require.NoError(t, err) + answer, _, err := Respond(init.Offer(), Binding{LocalWgPub: wgB, RemoteWgPub: wgA}) + require.NoError(t, err) + + id := ExchangeID{1, 2, 3, 4} + + offBytes, err := (&OfferMsg{ExchangeID: id, KEMOffer: init.Offer()}).Encode() + require.NoError(t, err) + typ, decoded, err := Decode(offBytes) + require.NoError(t, err) + require.Equal(t, MsgOffer, typ) + require.Equal(t, id, decoded.(*OfferMsg).ExchangeID) + require.Equal(t, init.Offer(), decoded.(*OfferMsg).KEMOffer) + + ansBytes, err := (&AnswerMsg{ExchangeID: id, KEMAnswer: answer}).Encode() + require.NoError(t, err) + typ, decoded, err = Decode(ansBytes) + require.NoError(t, err) + require.Equal(t, MsgAnswer, typ) + require.Equal(t, answer, decoded.(*AnswerMsg).KEMAnswer) + + confBytes := (&ConfirmMsg{ExchangeID: id}).Encode() + typ, decoded, err = Decode(confBytes) + require.NoError(t, err) + require.Equal(t, MsgConfirm, typ) + require.Equal(t, id, decoded.(*ConfirmMsg).ExchangeID) +} + +func TestDecodeRejects(t *testing.T) { + // too short + _, _, err := Decode([]byte{1, 1}) + require.Error(t, err) + + // wrong version + bad := make([]byte, headerSize+OfferSize) + bad[0] = byte(MsgOffer) + bad[1] = ProtocolVersion + 1 + _, _, err = Decode(bad) + require.Error(t, err) + + // unknown type + bad2 := make([]byte, headerSize) + bad2[0] = 99 + bad2[1] = ProtocolVersion + _, _, err = Decode(bad2) + require.Error(t, err) + + // offer with wrong payload size + _, err = (&OfferMsg{KEMOffer: []byte{1, 2, 3}}).Encode() + require.Error(t, err) +}