pnyx/packet.go

package pnyx

import (
	"bytes"
	"crypto/cipher"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/sha512"
	"encoding/binary"
	"io"
	"net"

	"fmt"
	"slices"

	"filippo.io/edwards25519"
)

type PacketType uint8
const (
  ID_LENGTH = 16
  PUBKEY_LENGTH = 32
  ECDH_LENGTH = 32
  SIGNATURE_LENGTH = 64
  HMAC_LENGTH = 64
  COMMAND_LENGTH = 1

  SESSION_OPEN_LENGTH = PUBKEY_LENGTH + ECDH_LENGTH + SIGNATURE_LENGTH
  SESSION_CONNECT_LENGTH = 2 + HMAC_LENGTH // + return addr string length

  SESSION_OPEN PacketType = iota
  SESSION_CONNECT
  SESSION_CLOSE
  SESSION_CLOSED
  SESSION_DATA
)

func ECDH(public ed25519.PublicKey, private ed25519.PrivateKey) ([]byte, error) {
  public_point, err := (&edwards25519.Point{}).SetBytes(public)
  if err != nil {
    return nil, err
  } 


  h := sha512.Sum512(private.Seed())
  private_scalar, err := (&edwards25519.Scalar{}).SetBytesWithClamping(h[:32])
  if err != nil {
    return nil, err
  }

  shared_point := public_point.ScalarMult(private_scalar, public_point)

  return shared_point.BytesMontgomery(), nil
}

func NewSessionOpen(key ed25519.PrivateKey) ([]byte, ed25519.PrivateKey, error) {
  if key == nil {
    return nil, nil, fmt.Errorf("Cannot create a SESSION_OPEN packet without a key")
  }

  public, private, err := ed25519.GenerateKey(rand.Reader)
  if err != nil {
    return nil, nil, fmt.Errorf("Failed to generate ecdh key: %w", err)
  }

  packet := make([]byte, COMMAND_LENGTH + SESSION_OPEN_LENGTH)
  cur := 0

  packet[0] = byte(SESSION_OPEN)
  cur += COMMAND_LENGTH

  copy(packet[cur:], []byte(key.Public().(ed25519.PublicKey)))
  cur += PUBKEY_LENGTH

  copy(packet[cur:], []byte(public)) 
  cur += PUBKEY_LENGTH

  signature := ed25519.Sign(key, packet[COMMAND_LENGTH:cur])
  copy(packet[cur:], signature)
  cur += SIGNATURE_LENGTH

  return packet, private, nil
}

func ParseSessionOpen(session_open []byte) (ed25519.PublicKey, ed25519.PublicKey, error) {
  if len(session_open) != SESSION_OPEN_LENGTH {
    return nil, nil, fmt.Errorf("Bad SESSION_OPEN length: %d/%d", len(session_open), SESSION_OPEN_LENGTH)
  }

  cur := 0

  client_pubkey := (ed25519.PublicKey)(session_open[cur:cur+PUBKEY_LENGTH])
  cur += PUBKEY_LENGTH

  client_ecdh := (ed25519.PublicKey)(session_open[cur:cur+PUBKEY_LENGTH])
  cur += PUBKEY_LENGTH
  
  digest := session_open[:cur]
  signature := session_open[cur:cur+SIGNATURE_LENGTH]
  cur += SIGNATURE_LENGTH

  if ed25519.Verify(client_pubkey, digest, signature) == false {
    return nil, nil, fmt.Errorf("SESSION_OPEN signature verification failed")
  }

  return client_pubkey, client_ecdh, nil
}

func NewSessionConnect(address *net.UDPAddr, session_secret []byte) []byte {
  packet := make([]byte, COMMAND_LENGTH + ID_LENGTH + SESSION_CONNECT_LENGTH + len(address.String()))
  cur := 0

  packet[cur] = byte(SESSION_CONNECT)
  cur += COMMAND_LENGTH

  session_id := [16]byte(ID[SessionID](session_secret))
  copy(packet[cur:], session_id[:])
  cur += ID_LENGTH

  binary.BigEndian.PutUint16(packet[cur:], uint16(len(address.String())))
  cur += 2
  
  copy(packet[cur:], []byte(address.String()))
  cur += len(address.String())

  hmac := sha512.Sum512(append(packet[COMMAND_LENGTH+ID_LENGTH:cur], session_secret...))
  copy(packet[cur:], hmac[:])

  return packet
}

func ParseSessionConnect(session_connect []byte, session_secret []byte) (*net.UDPAddr, error) {
  if len(session_connect) < SESSION_CONNECT_LENGTH {
    return nil, fmt.Errorf("Bad session connect length: %d/%d", len(session_connect), SESSION_CONNECT_LENGTH)
  }

  cur := 0

  address_length := int(binary.BigEndian.Uint16(session_connect[cur:cur+2]))
  cur += 2

  if len(session_connect) != (SESSION_CONNECT_LENGTH + address_length) {
    return nil, fmt.Errorf("Bad session connect length: %d/%d", len(session_connect), SESSION_CONNECT_LENGTH + address_length)
  }

  address := string(session_connect[cur:cur+address_length])
  cur += address_length

  hmac_digest := make([]byte, cur)
  copy(hmac_digest, session_connect[:cur])

  hmac := session_connect[cur:cur+HMAC_LENGTH]
  cur += HMAC_LENGTH

  calculated_hmac := sha512.Sum512(append(hmac_digest, session_secret...))
  if slices.Compare(hmac, calculated_hmac[:]) != 0 {
    return nil, fmt.Errorf("Session connect bad HMAC")
  }

  addr, err := net.ResolveUDPAddr("udp", address)
  if err != nil {
    return nil, fmt.Errorf("Error parsing return address: %w", err)
  }

  return addr, nil
}

func NewSessionData(session *Session, packet []byte) ([]byte, error) {
  iv := make([]byte, 32)
  for i := 0; i < 4; i++ {
    binary.BigEndian.PutUint64(iv[i*8:], session.iv_generator.Uint64())
  }

  stream := cipher.NewOFB(session.cipher, iv[:])
  header := make([]byte, COMMAND_LENGTH + ID_LENGTH)
  header[0] = byte(SESSION_DATA)
  copy(header[1:], session.ID[:])

  packet_encrypted := bytes.NewBuffer(header)
  writer := &cipher.StreamWriter{S: stream, W: packet_encrypted}

  _, err := io.Copy(writer, bytes.NewBuffer(packet))
  if err != nil {
    return nil, err
  }

  return packet_encrypted.Bytes(), nil
}

func ParseSessionData(session *Session, data []byte) ([]byte, error) {
  iv := data[0:32]

  stream := cipher.NewOFB(session.cipher, iv)
  var packet_clear bytes.Buffer
  reader := &cipher.StreamReader{S: stream, R: bytes.NewBuffer(data)}
  _, err := io.Copy(&packet_clear, reader)
  if err != nil {
    return nil, err
  }

  return packet_clear.Bytes(), nil
}