wg_init.go

package floaties

import (
	"encoding/binary"
	"fmt"
	"net"
	"time"
)

// msg = handshake_initiation {
//	 u8 message_type									0
//	 u8 reserved_zero[3]									1
//	 u32 sender_index									4
//	 u8 unencrypted_ephemeral[32]							8
//	 u8 encrypted_static[AEAD_LEN(32)]				40
//	 u8 encrypted_timestamp[AEAD_LEN(12)]			88
//	 u8 mac1[16]										116
//	 u8 mac2[16]										132
// }

// gross stream-of-consciousness protocol message layout for
// the WireGuard initiation message

const (
	InitSize           = 148
	InitOffType        = 0
	InitOffZero        = 1
	InitOffSenderIndex = 4
	InitOffEphemeral   = 8
	InitOffStatic      = 40
	InitOffTimestamp   = 88
	InitOffMac1        = 116
	InitOffMac2        = 132
)

// WgInitiator holds the keying information for a WireGuard session
type WgInitiator struct {
	StaticPriv  [32]byte
	StaticPub   [32]byte
	EphemPriv   [32]byte
	EphemPub    [32]byte
	RStaticPub  [32]byte
	SenderIndex uint32
}

// NewInitiator takes the local static private key and the remote static public
// key, generates a new local ephemeral key, and returns the keying
// information
func NewInitiator(priv, pub [32]byte, index uint32) WgInitiator {
	eph := DH_GENERATE()

	return WgInitiator{
		StaticPriv:  priv,
		StaticPub:   DH_PUBKEY(priv),
		EphemPriv:   eph,
		EphemPub:    DH_PUBKEY(eph),
		RStaticPub:  pub,
		SenderIndex: index,
	}
}

// Initiation generates a WireGuard protocol initiation message,
func (i *WgInitiator) Initiation() ([]byte, error) {
	// this is objectively horrible, copy-as-much-as-you-can
	// potato crypto code, and you should never do anything serious
	// with it.

	msg := make([]byte, InitSize)

	// the chaining key generates the ChaPoly AEAD cipher key and
	// incorporates the DH shared secrets generated by the handshake
	chaining_key := HASH([]byte(CONSTRUCTION))

	// the hash tracks the transcript of the values sent in the
	// handshake and is the additional data in invocations of the
	// ChaPoly AEAD
	hash := HASH(concat(HASH(concat(chaining_key, []byte(IDENTIFIER))), i.RStaticPub[:]))

	msg[0] = 1 // initiation message
	binary.LittleEndian.PutUint32(msg[4:], i.SenderIndex)

	// hash and chaining key include our local ephemeral public key

	copy(msg[InitOffEphemeral:], i.EphemPub[:])
	hash = HASH(concat(hash, i.EphemPub[:]))

	temp := HMAC(chaining_key, i.EphemPub[:])
	chaining_key = HMAC(temp, []byte{1})

	// DH ephemeral-static, into the chaining key

	dhv, err := DH(i.EphemPriv, i.RStaticPub)
	if err != nil {
		return nil, fmt.Errorf("initiator: es dh: %w", err)
	}
	temp = HMAC(chaining_key, dhv)
	chaining_key = HMAC(temp, []byte{1})
	key := HMAC(temp, concat(chaining_key, []byte{2}))

	// record our encrypted local static (identity) public key into
	// the packet, and hash the ciphertext into the transcript hash

	copy(msg[InitOffStatic:], AEAD(key, 0, i.StaticPub[:], hash))

	hash = HASH(concat(hash, msg[InitOffStatic:InitOffStatic+AEAD_LEN(32)]))

	// DH static-static, into the chaining key

	dhv, err = DH(i.StaticPriv, i.RStaticPub)
	if err != nil {
		return nil, fmt.Errorf("initiator: ss dh: %w", err)
	}
	temp = HMAC(chaining_key, dhv)
	chaining_key = HMAC(temp, []byte{1})
	key = HMAC(temp, concat(chaining_key, []byte{2}))

	// record an encrypted timestamp into the packet, and hash the
	// ciphertext into the transcript hash

	copy(msg[InitOffTimestamp:], AEAD(key, 0, TAI64N(), hash))

	hash = HASH(concat(hash, msg[InitOffTimestamp:InitOffTimestamp+AEAD_LEN(12)]))

	// MAC the whole packet with keyed blake2 as a DDOS defense

	copy(msg[InitOffMac1:], MAC(HASH(concat([]byte(LABEL_MAC1), i.RStaticPub[:])), msg[:InitOffMac1]))

	return msg, nil
}

// CheckInitiate generates an initiation message, sends it to `addr`, and
// waits for a response; we do nothing with the response except generate a
// non-error return when we see it, since we've verified we can talk to
// WireGuard
func (i *WgInitiator) CheckInitiate(addr string) error {
	raddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		return fmt.Errorf("initiate: %w", err)
	}

	msg, err := i.Initiation()
	if err != nil {
		return fmt.Errorf("initiate: generate msg: %w", err)
	}

	uc, err := net.DialUDP("udp", nil, raddr)
	if err != nil {
		return fmt.Errorf("initiate: %w", err)
	}

	_, err = uc.Write(msg)
	if err != nil {
		return fmt.Errorf("initiate: %w", err)
	}

	uc.SetDeadline(time.Now().Add(2 * time.Second))
	rbuf := make([]byte, 1500)
	_, err = uc.Read(rbuf)
	if err != nil {
		return fmt.Errorf("initiate: %w", err)
	}

	return nil
}

func RecoverIdentityFromInitiation(buf []byte, priv, pub [32]byte) (pubk []byte, err error) {
	if len(buf) < 132 {
		return nil, fmt.Errorf("decrypt initiation: too short")
	}

	if buf[0] != 1 {
		return nil, fmt.Errorf("decrypt initiation: not an initiation message")
	}

	chaining_key := HASH([]byte(CONSTRUCTION))
	hash := HASH(concat(HASH(concat(chaining_key, []byte(IDENTIFIER))), pub[:]))

	var ephem [32]byte
	copy(ephem[:], buf[InitOffEphemeral:InitOffEphemeral+32])
	hash = HASH(concat(hash, ephem[:]))

	temp := HMAC(chaining_key, ephem[:])
	chaining_key = HMAC(temp, []byte{1})

	dhv, err := DH(priv, ephem)
	if err != nil {
		return nil, fmt.Errorf("decrypt initiation: %w", err)
	}

	temp = HMAC(chaining_key, dhv)
	chaining_key = HMAC(temp, []byte{1})
	key := HMAC(temp, concat(chaining_key, []byte{2}))

	clientPubEncrypted := make([]byte, AEAD_LEN(32))

	copy(clientPubEncrypted, buf[InitOffStatic:InitOffStatic+AEAD_LEN(32)])
	clientPub, err := AEAD_DECRYPT(key, 0, clientPubEncrypted, hash)
	if err != nil {
		return nil, fmt.Errorf("decrypt initiation: %w", err)
	}

	return clientPub, nil
}