pnyx/client.go

package pnyx

import (
	"crypto/ed25519"
	"crypto/rand"
	"fmt"
	"net"
	"sync/atomic"
	"time"
)

type ClientState uint8

const (
  CLIENT_SESSION_CONNECTING ClientState = iota
  CLIENT_SESSION_CONNECTED
)

type Client struct {
  Key ed25519.PrivateKey
  Session Session
  remote string

  connection *net.UDPConn

  to_send chan []byte
  data_fn func(Payload)error
  state ClientState
  active atomic.Bool
}

func(client *Client) Remote() string {
  return client.remote
}

func(client *Client) Active() bool {
  return client.active.Load()
}

const CLIENT_UDP_BUFFER = 2048
const CLIENT_MAX_CONNECT_ATTEMPTS = 5

func(client *Client) Log(format string, vals ...any) {
  fmt.Printf("%s\n", fmt.Sprintf(format, vals...))
}

func(client *Client) listen_udp() {
  buffer := [CLIENT_UDP_BUFFER]byte{}
  for client.active.Load() {
    read, err := client.connection.Read(buffer[:])
    if err != nil {
      client.Log("Client listen error - %s", err)
    } else if read == 0 {
      client.Log("Client listen error - no data in packet")
    } else {
      packet_type := SessionPacketType(buffer[0])
      switch client.state {
      case CLIENT_SESSION_CONNECTING:
        switch packet_type {
        case SESSION_CONNECTED:
          client.state = CLIENT_SESSION_CONNECTED
        case SESSION_CLOSED:
          client.Log("Server repsonded to session connect with SESSION_CLOSED")
          client.active.Store(false)
        default:
          client.Log("Bad session packet type 0x%02x for client in state SESSION_CONNECTING", packet_type)
        }
      case CLIENT_SESSION_CONNECTED:
        switch packet_type {
        case SESSION_DATA:
          if len(buffer) < SESSION_ID_LENGTH + 1 {
            client.Log("Not enough data to decode SESSION_DATA packet %d/%d", len(buffer), SESSION_ID_LENGTH + 1)
            continue
          }

          session_id := SessionID(buffer[1:1+SESSION_ID_LENGTH])
          if session_id != client.Session.ID {
            client.Log("Session ID of data packet does not match client.Session %s =/= %s", session_id, client.Session.ID)
            continue
          }

          data, err := ParseSessionData(&client.Session, buffer[1+SESSION_ID_LENGTH:read])
          if err != nil {
            client.Log("Error parsing session data: %s", err)
            continue
          }

          payload, err := ParsePacket(data)
          if err != nil {
            client.Log("Error parsing packet from session data: %s", err)
            continue
          }

          switch payload.(type) {
          case PingPacket:
            err = client.Send(NewPingPacket())
            if err != nil {
              client.Log("Error sending ping packet: %s", err)
            }
          }

          err = client.data_fn(payload)
          if err != nil {
            client.Log("Error running data_fn: %s", err)
          }
        default:
          client.Log("Bad session packet type 0x%02x for client in state SESSION_CONNECTED", packet_type)
        }
      default:
        client.Log("Received packet while in bad client state 0x%02x", client.state)
      }
    }
  }
}

func NewClient(key ed25519.PrivateKey, remote string, data_fn func(Payload)error) (*Client, error) {
  if key == nil {
    return nil, fmt.Errorf("Need a key to create a client, passed nil")
  } else if data_fn == nil {
    return nil, fmt.Errorf("Need a function to run with session data")
  }

  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
  }

  var response = [512]byte{}
  var session Session
  for attempts := 0; attempts <= CLIENT_MAX_CONNECT_ATTEMPTS; attempts++ {
    _, err = connection.Write(session_open)
    if err != nil {
      return nil, err
    }

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

    if response[0] != byte(SESSION_OPENED) {
      return nil, fmt.Errorf("Invalid SESSION_OPEN response: %x", response[0])
    }

    session, err = ParseSessionOpened(nil, ecdh_private, response[COMMAND_LENGTH:read])
    if err == nil {
      break
    }

    net_error, ok := err.(net.Error)
    if ok == false {
      return nil, err
    } else if net_error.Timeout() {
      if attempts == CLIENT_MAX_CONNECT_ATTEMPTS {
        return nil, fmt.Errorf("Failed to connect to server at %s", remote)
      }
    } else {
      return nil, err
    }
  }

  session_connect := NewSessionTimed(SESSION_CONNECT, key, &session, time.Now())
  _, err = connection.Write(session_connect)
  if err != nil {
    return nil, err
  }
 
  client := &Client{
    Key: key,
    Session: session,
    remote: remote,

    connection: connection,
    to_send: make(chan []byte, 1000),
    data_fn: data_fn,
    state: CLIENT_SESSION_CONNECTING,
  }

  client.active.Store(true)

  go client.listen_udp()
  go client.send_queue()

  return client, nil
}

func(client *Client) send_queue() {
  for client.active.Load() {
    packet := <- client.to_send
    if packet == nil {
      break
    }
    _, err := client.connection.Write(packet)
    if err != nil {
      client.Log("Send error: %s", err)
    }
  }
}

func(client *Client) Send(packet *Packet) error {
  data, err := packet.MarshalBinary()
  if err != nil {
    return err
  }

  wrapped, err := NewSessionData(&client.Session, data)
  if err != nil {
    return err
  }

  select {
  case client.to_send <- wrapped:
    return nil
  default:
    return fmt.Errorf("Channel overflow")
  }
}

func(client *Client) Close() error {
  var err error = fmt.Errorf("Client not active")
  if client.active.CompareAndSwap(true, false) {
    close(client.to_send)
    client.connection.Write(NewSessionTimed(SESSION_CLOSE, client.Key, &client.Session, time.Now()))
    err = client.connection.Close()
  }
  return err
}