pnyx/client.go

package pnyx

import (
	"crypto/ed25519"
	"crypto/rand"
  mrand "math/rand"
  "encoding/binary"
	"crypto/sha512"
  "crypto/aes"
	"net"
	"sync"
  "fmt"

	"github.com/google/uuid"
)

type ClientID uuid.UUID
func(id ClientID) String() string {
  return uuid.UUID(id).String()
}

type ClientState uint8

const (
  CLIENT_SESSION_CREATE ClientState = iota
  CLIENT_SESSION_CONNECT
  CLIENT_SESSION_CONNECTED
)

type Client struct {
  Key ed25519.PrivateKey
  Session Session
  ConnectionLock sync.Mutex
  Connection *net.UDPConn
}

func ID[T ~[16]byte, V ~[]byte](data V) T { 
  hash := sha512.Sum512(data)
  return (T)(hash[0:16])
}

func NewClient(key ed25519.PrivateKey, remote string) (*Client, error) {
  if key == nil {
    var err error
    _, key, err = ed25519.GenerateKey(rand.Reader)
    if err != nil {
      return nil, err
    }
  }

  seed_bytes := make([]byte, 8)
  read, err := rand.Read(seed_bytes)
  if err != nil {
    return nil, err
  } else if read != 8 {
    return nil, fmt.Errorf("Failed to create IV seed for client")
  }

  address, err := net.ResolveUDPAddr("udp", remote)
  if err != nil {
    return nil, err
  }

  connection, err := net.DialUDP("udp", nil, address)
  if err != nil {
    return nil, err
  }

  session_open, ecdh_private, err := NewSessionOpen(key)
  if err != nil {
    connection.Close()
    return nil, err
  }

  _, err = connection.Write(session_open)
  if err != nil {
    return nil, err
  }

  var response = [512]byte{}
  read, _, err = connection.ReadFromUDP(response[:])
  if err != nil {
    return nil, err
  }
  
  if response[0] != byte(SESSION_OPEN) {
    return nil, fmt.Errorf("Invalid SESSION_OPEN response: %x", response[0])
  }

  server_pubkey, ecdh_public, err := ParseSessionOpen(response[COMMAND_LENGTH:read])
  if err != nil {
    return nil, err
  }

  secret, err := ECDH(ecdh_public, ecdh_private)
  if err != nil {
    return nil, err
  }

  session_cipher, err := aes.NewCipher(secret)
  if err != nil {
    return nil, err
  }

  session_connect := NewSessionConnect(connection.LocalAddr().(*net.UDPAddr), secret)
  _, err = connection.Write(session_connect)
  if err != nil {
    return nil, err
  }

  read, _, err = connection.ReadFromUDP(response[:])
  if err != nil {
    return nil, err
  }

  return &Client{
    Key: key,
    Session: Session{
      ID: ID[SessionID](secret),
      remote: address,
      Peer: ID[ClientID](server_pubkey),
      secret: secret,
      cipher: session_cipher,
      iv_generator: mrand.NewSource(int64(binary.BigEndian.Uint64(seed_bytes))).(mrand.Source64),
    },
    Connection: connection,
  }, nil
}