pnyx/client.go

package pnyx

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

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 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_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 {
    return nil, err
  }

  session_connect := NewSessionTimed(SESSION_CONNECT, key, &session, time.Now())
  _, 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,
    Connection: connection,
  }, nil
}

func(client *Client) Send(packet *Packet) error {
  client.ConnectionLock.Lock()
  defer client.ConnectionLock.Unlock()

  if client.Connection == nil {
    return fmt.Errorf("Client is not connected")
  }

  data, err := packet.MarshalBinary()
  if err != nil {
    return err
  }

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

  _, err = client.Connection.Write(wrapped)
  return err
}

func(client *Client) Close() error {
  client.ConnectionLock.Lock()
  defer client.ConnectionLock.Unlock()

  if client.Connection == nil {
    return fmt.Errorf("No connection to close")
  }

  client.Connection.Write(NewSessionTimed(SESSION_CLOSE, client.Key, &client.Session, time.Now()))
  err := client.Connection.Close()
  client.Connection = nil

  return err
}