pnyx/client.go

package pnyx

import (
	"crypto/ed25519"
	"crypto/rand"
	"crypto/sha512"
	"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
  ConnectionLock sync.Mutex
  Connection *net.UDPConn
  State ClientState
}

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

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

  return Client{
    Key: key,
    State: CLIENT_SESSION_CREATE,
  }, nil
}

func(client *Client) Connect(remote string) (ed25519.PublicKey, []byte, error) {
  client.ConnectionLock.Lock()
  defer client.ConnectionLock.Unlock()
  
  address, err := net.ResolveUDPAddr("udp", remote)
  if err != nil {
    return nil, nil, err
  }

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

  session_open, ecdh_private, err := NewSessionOpen(client.Key)
  if err != nil {
    client.Connection.Close()
    client.Connection = nil
    return nil, nil, err
  }

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

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

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

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

  client.State = CLIENT_SESSION_CONNECT

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

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

  return server_public, secret, nil
}