qcc/mobile/client.go

package mobile

import (
	"context"
	"crypto/ed25519"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"crypto/tls"
	"crypto/x509"
	"encoding/binary"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	"github.com/apernet/quic-go"
	"github.com/niko/qcc/internal/auth"
	"github.com/niko/qcc/internal/protocol"
	"github.com/niko/qcc/pkg/types"
	"golang.org/x/crypto/chacha20poly1305"
	"golang.org/x/crypto/curve25519"
	"golang.org/x/crypto/hkdf"
)

type Client struct {
	mu        sync.Mutex
	conn      *quic.Conn
	transport *quic.Transport
	udpConn   net.PacketConn

	privKey   ed25519.PrivateKey
	pubKey    ed25519.PublicKey

	number    string
	certDER   []byte
	caCertDER []byte

	serverAddr string

	callbacks Callbacks

	ctx    context.Context
	cancel context.CancelFunc

	e2eePriv  [32]byte
	e2eePub   [32]byte
	mediaKey  []byte
	activeCall bool
	callID    uint64

	authOpened bool
}

type Callbacks interface {
	OnIncomingCall(number string)
	OnCallAccepted()
	OnCallEnded()
	OnCallRejected()
	OnError(msg string)
	OnMedia(data []byte)
}

func NewClient() *Client {
	pub, priv, _ := ed25519.GenerateKey(rand.Reader)
	return &Client{
		privKey: priv,
		pubKey:  pub,
	}
}

func NewClientWithKey(privKey []byte) *Client {
	priv := ed25519.PrivateKey(privKey)
	pub := priv.Public().(ed25519.PublicKey)
	return &Client{
		privKey: priv,
		pubKey:  pub,
	}
}

func Ping() string {
	return "pong"
}

func (c *Client) PrivateKey() []byte {
	c.mu.Lock()
	defer c.mu.Unlock()
	return []byte(c.privKey)
}

func (c *Client) PublicKey() []byte {
	c.mu.Lock()
	defer c.mu.Unlock()
	return []byte(c.pubKey)
}

func (c *Client) SetCallbacks(cbs Callbacks) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.callbacks = cbs
}

func (c *Client) CertDER() []byte {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.certDER
}

func (c *Client) CACertDER() []byte {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.caCertDER
}

func (c *Client) SetCertDER(der []byte) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.certDER = der
}

func (c *Client) SetCACertDER(der []byte) {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.caCertDER = der
}

func (c *Client) AuthenticateB64(certB64 string, caCertB64 string) (number string, err error) {
	defer func() {
		if r := recover(); r != nil {
			number = ""
			err = fmt.Errorf("panic in AuthenticateB64: %v", r)
		}
	}()

	var certDER []byte
	if certB64 != "" {
		certDER, err = base64.StdEncoding.DecodeString(certB64)
		if err != nil {
			return "", fmt.Errorf("decode cert: %w", err)
		}
	}

	c.mu.Lock()
	conn := c.conn
	c.mu.Unlock()

	if conn == nil {
		return "", fmt.Errorf("not connected")
	}

	authSt, streamErr := conn.OpenStream()
	if streamErr != nil {
		return "", streamErr
	}

	if streamErr = protocol.WriteFrame(authSt, byte(types.OpGetChallenge), nil); streamErr != nil {
		authSt.Close()
		return "", fmt.Errorf("send challenge: %w", streamErr)
	}

	frame, streamErr := protocol.ReadFrame(authSt)
	if streamErr != nil {
		authSt.Close()
		return "", fmt.Errorf("read challenge: %w", streamErr)
	}

	if types.OpCode(frame.OpCode) != types.OpChallenge {
		authSt.Close()
		return "", fmt.Errorf("expected challenge, got op=%x", frame.OpCode)
	}

	challenge := auth.UnmarshalChallenge(frame.Payload)
	if challenge == nil {
		authSt.Close()
		return "", fmt.Errorf("invalid challenge payload")
	}

	if len(certDER) > 0 {
		c.mu.Lock()
		privKey := c.privKey
		c.mu.Unlock()

		challengeData := challenge.Marshal()
		signature := ed25519.Sign(privKey, challengeData)

		payload := make([]byte, 1+64+2+len(certDER))
		payload[0] = 64
		copy(payload[1:], signature)
		binary.BigEndian.PutUint16(payload[1+64:], uint16(len(certDER)))
		copy(payload[1+64+2:], certDER)

		if streamErr = protocol.WriteFrame(authSt, byte(types.OpAuthCert), payload); streamErr != nil {
			authSt.Close()
			return "", streamErr
		}

		frame, streamErr = protocol.ReadFrame(authSt)
		if streamErr != nil {
			authSt.Close()
			return "", streamErr
		}

		if types.OpCode(frame.OpCode) == types.OpError {
			authSt.Close()
			errMsg := parseError(frame.Payload)
			return "", fmt.Errorf("server error: %s", errMsg)
		}

		if types.OpCode(frame.OpCode) != types.OpIdentity {
			authSt.Close()
			return "", fmt.Errorf("expected identity, got op=%x", frame.OpCode)
		}

		number, _, _ = parseIdentityPayload(frame.Payload)
		if number == "" {
			authSt.Close()
			return "", fmt.Errorf("invalid identity payload")
		}

		c.mu.Lock()
		c.number = number
		c.authOpened = true
		c.mu.Unlock()
		go c.handleIncomingStreams()
		return number, nil
	}

	solution := c.solvePoW(challenge)
	if solution == nil {
		authSt.Close()
		return "", fmt.Errorf("PoW solve failed")
	}

	solnPayload := make([]byte, 40)
	copy(solnPayload[0:8], solution.ClientNonce[:])
	copy(solnPayload[8:40], solution.ClientPubKey[:])

	if streamErr = protocol.WriteFrame(authSt, byte(types.OpSolve), solnPayload); streamErr != nil {
		authSt.Close()
		return "", fmt.Errorf("send solution: %w", streamErr)
	}

	frame, streamErr = protocol.ReadFrame(authSt)
	if streamErr != nil {
		authSt.Close()
		return "", fmt.Errorf("read response: %w", streamErr)
	}

	if types.OpCode(frame.OpCode) == types.OpError {
		authSt.Close()
		errMsg := parseError(frame.Payload)
		return "", fmt.Errorf("server error: %s", errMsg)
	}

	if types.OpCode(frame.OpCode) != types.OpIdentity {
		authSt.Close()
		return "", fmt.Errorf("expected identity, got op=%x", frame.OpCode)
	}

	number, newCertDER, newCACertDER := parseIdentityPayload(frame.Payload)
	if number == "" {
		authSt.Close()
		return "", fmt.Errorf("invalid identity payload")
	}

	c.mu.Lock()
	c.number = number
	c.certDER = newCertDER
	c.caCertDER = newCACertDER
	c.authOpened = true
	c.mu.Unlock()

	go c.handleIncomingStreams()

	return number, nil
}

func (c *Client) Connect(addr string) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	c.serverAddr = addr

	host, _, err := net.SplitHostPort(addr)
	if err != nil {
		host = addr
	}

	tlsConf := &tls.Config{
		ServerName:         host,
		InsecureSkipVerify: true,
		NextProtos:         []string{"qcc"},
		MinVersion:         tls.VersionTLS13,
	}

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

	udpConn, err := net.ListenUDP("udp", nil)
	if err != nil {
		return fmt.Errorf("listen udp: %w", err)
	}
	c.udpConn = udpConn

	tr := &quic.Transport{Conn: udpConn}
	c.transport = tr

	ctx := context.Background()
	conn, err := tr.Dial(ctx, udpAddr, tlsConf, &quic.Config{
		EnableDatagrams: true,
		KeepAlivePeriod: 15 * time.Second,
	})
	if err != nil {
		udpConn.Close()
		c.transport = nil
		c.udpConn = nil
		return fmt.Errorf("quic dial: %w", err)
	}

	c.conn = conn
	c.ctx, c.cancel = context.WithCancel(ctx)

	return nil
}

func (c *Client) Authenticate() (number string, err error) {
	defer func() {
		if r := recover(); r != nil {
			number = ""
			err = fmt.Errorf("panic in Authenticate: %v", r)
		}
	}()

	c.mu.Lock()
	conn := c.conn
	c.mu.Unlock()

	if conn == nil {
		return "", fmt.Errorf("not connected")
	}

	authSt, streamErr := conn.OpenStream()
	if streamErr != nil {
		return "", streamErr
	}

	if streamErr = protocol.WriteFrame(authSt, byte(types.OpGetChallenge), nil); streamErr != nil {
		authSt.Close()
		return "", fmt.Errorf("send challenge: %w", streamErr)
	}

	frame, streamErr := protocol.ReadFrame(authSt)
	if streamErr != nil {
		authSt.Close()
		return "", fmt.Errorf("read challenge: %w", streamErr)
	}

	if types.OpCode(frame.OpCode) != types.OpChallenge {
		authSt.Close()
		return "", fmt.Errorf("expected challenge, got op=%x", frame.OpCode)
	}

	challenge := auth.UnmarshalChallenge(frame.Payload)
	if challenge == nil {
		authSt.Close()
		return "", fmt.Errorf("invalid challenge payload")
	}

	solution := c.solvePoW(challenge)
	if solution == nil {
		authSt.Close()
		return "", fmt.Errorf("PoW solve failed")
	}

	solnPayload := make([]byte, 40)
	copy(solnPayload[0:8], solution.ClientNonce[:])
	copy(solnPayload[8:40], solution.ClientPubKey[:])

	if streamErr = protocol.WriteFrame(authSt, byte(types.OpSolve), solnPayload); streamErr != nil {
		authSt.Close()
		return "", fmt.Errorf("send solution: %w", streamErr)
	}

	frame, streamErr = protocol.ReadFrame(authSt)
	if streamErr != nil {
		authSt.Close()
		return "", fmt.Errorf("read response: %w", streamErr)
	}

	if types.OpCode(frame.OpCode) == types.OpError {
		authSt.Close()
		errMsg := parseError(frame.Payload)
		return "", fmt.Errorf("server error: %s", errMsg)
	}

	if types.OpCode(frame.OpCode) != types.OpIdentity {
		authSt.Close()
		return "", fmt.Errorf("expected identity, got op=%x", frame.OpCode)
	}

	number, newCertDER, newCACertDER := parseIdentityPayload(frame.Payload)
	if number == "" {
		authSt.Close()
		return "", fmt.Errorf("invalid identity payload")
	}

	c.mu.Lock()
	c.number = number
	c.certDER = newCertDER
	c.caCertDER = newCACertDER
	c.authOpened = true
	c.mu.Unlock()

	go c.handleIncomingStreams()

	return number, nil
}

func (c *Client) solvePoW(challenge *auth.PoWChallenge) *auth.PoWSolution {
	bits := challenge.Bits
	var nonceCounter uint64

	var solution auth.PoWSolution
	copy(solution.ClientPubKey[:], c.pubKey)

	hashInput := make([]byte, 32)
	binary.BigEndian.PutUint64(hashInput[0:8], uint64(challenge.Timestamp))
	copy(hashInput[8:16], challenge.ServerNonce[:])

	for {
		binary.BigEndian.PutUint64(hashInput[16:24], nonceCounter)
		copy(hashInput[24:], solution.ClientPubKey[:])

		hash := sha256.Sum256(hashInput)

		leadingZeros := 0
		for i := 0; i < 32; i++ {
			if hash[i] == 0 {
				leadingZeros += 8
			} else {
				for b := byte(0x80); b != 0; b >>= 1 {
					if hash[i]&b == 0 {
						leadingZeros++
					} else {
						goto checkDone
					}
				}
			}
		}
	checkDone:
		if leadingZeros >= bits {
			binary.BigEndian.PutUint64(solution.ClientNonce[:], nonceCounter)
			return &solution
		}
		nonceCounter++
	}
}

func (c *Client) MyNumber() string {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.number
}

func (c *Client) ServerAddr() string {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.serverAddr
}

func (c *Client) IsConnected() bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.conn != nil
}

func (c *Client) Disconnect() {
	c.mu.Lock()
	defer c.mu.Unlock()

	if c.cancel != nil {
		c.cancel()
	}
	if c.conn != nil {
		c.conn.CloseWithError(0, "client disconnect")
	}
	if c.transport != nil {
		c.transport.Close()
	}
	if c.udpConn != nil {
		c.udpConn.Close()
	}
	c.conn = nil
	c.transport = nil
	c.udpConn = nil
	c.number = ""
	c.activeCall = false
	c.callID = 0
	c.mediaKey = nil
	c.authOpened = false
}

func (c *Client) handleIncomingStreams() {
	defer func() { recover() }()

	c.mu.Lock()
	conn := c.conn
	ctx := c.ctx
	c.mu.Unlock()

	if conn == nil {
		return
	}

	for {
		st, err := conn.AcceptStream(ctx)
		if err != nil {
			return
		}
		go func() {
			defer func() { recover() }()
			c.handleStream(st)
		}()
	}
}

func (c *Client) handleStream(st *quic.Stream) {
	defer st.Close()
	defer func() { recover() }()

	for {
		frame, err := protocol.ReadFrame(st)
		if err != nil {
			return
		}

		switch types.OpCode(frame.OpCode) {
		case types.OpRing:
			number, callerE2EE := parseRingPayload(frame.Payload)
			c.mu.Lock()
			c.e2eePub = callerE2EE
			c.mu.Unlock()
			if c.callbacks != nil {
				c.callbacks.OnIncomingCall(number)
			}

		case types.OpPeerAccept:
			var peerE2EE [32]byte
			if len(frame.Payload) >= 32 {
				copy(peerE2EE[:], frame.Payload[:32])
			}
			c.mu.Lock()
			c.e2eePub = peerE2EE
			c.deriveMediaKey()
			c.activeCall = true
			c.mu.Unlock()
			if c.callbacks != nil {
				c.callbacks.OnCallAccepted()
			}

		case types.OpPeerEnd:
			c.mu.Lock()
			c.activeCall = false
			c.callID = 0
			c.mediaKey = nil
			c.mu.Unlock()
			if c.callbacks != nil {
				c.callbacks.OnCallEnded()
			}

		case types.OpError:
			errMsg := parseError(frame.Payload)
			if c.callbacks != nil {
				c.callbacks.OnError(errMsg)
			}
			return
		}
	}
}

func (c *Client) handleDatagrams() {
	defer func() { recover() }()

	c.mu.Lock()
	conn := c.conn
	ctx := c.ctx
	c.mu.Unlock()

	if conn == nil {
		return
	}

	for {
		data, err := conn.ReceiveDatagram(ctx)
		if err != nil {
			return
		}

		pkt, ok := protocol.UnmarshalMediaPacket(data)
		if !ok {
			continue
		}

		c.mu.Lock()
		cid := c.callID
		key := c.mediaKey
		c.mu.Unlock()

		if pkt.CallID != cid {
			continue
		}

		if key != nil && len(pkt.Payload) > 0 {
			decrypted, err := decryptE2EE(key, pkt.Payload)
			if err == nil {
				if c.callbacks != nil {
					c.callbacks.OnMedia(decrypted)
				}
			}
		}
	}
}

func (c *Client) Dial(number string) error {
	c.mu.Lock()
	conn := c.conn
	c.mu.Unlock()
	if conn == nil {
		return fmt.Errorf("not connected")
	}

	e2eeKP, err := generateE2EEKey()
	if err != nil {
		return fmt.Errorf("generate e2ee key: %w", err)
	}

	c.mu.Lock()
	c.e2eePriv = e2eeKP.PrivateKey
	c.mu.Unlock()

	payload := make([]byte, 1+len(number)+32)
	payload[0] = byte(len(number))
	copy(payload[1:], number)
	copy(payload[1+len(number):], e2eeKP.PublicKey[:])

	st, err := conn.OpenStream()
	if err != nil {
		return fmt.Errorf("open call stream: %w", err)
	}

	if err := protocol.WriteFrame(st, byte(types.OpDial), payload); err != nil {
		st.Close()
		return fmt.Errorf("send dial: %w", err)
	}

	go func() {
		defer st.Close()
		c.handleStream(st)
	}()

	return nil
}

func (c *Client) AcceptCall() error {
	e2eeKP, err := generateE2EEKey()
	if err != nil {
		return err
	}

	c.mu.Lock()
	c.e2eePriv = e2eeKP.PrivateKey
	c.mu.Unlock()

	payload := make([]byte, 32)
	copy(payload, e2eeKP.PublicKey[:])

	c.mu.Lock()
	conn := c.conn
	c.mu.Unlock()
	if conn == nil {
		return fmt.Errorf("not connected")
	}

	st, err := conn.OpenStream()
	if err != nil {
		return fmt.Errorf("open stream: %w", err)
	}

	if err := protocol.WriteFrame(st, byte(types.OpAccept), payload); err != nil {
		st.Close()
		return fmt.Errorf("send accept: %w", err)
	}

	go func() {
		defer st.Close()
		c.handleStream(st)
	}()

	return nil
}

func (c *Client) RejectCall() error {
	c.mu.Lock()
	conn := c.conn
	c.mu.Unlock()
	if conn == nil {
		return nil
	}

	st, err := conn.OpenStream()
	if err != nil {
		return err
	}
	defer st.Close()

	return protocol.WriteFrame(st, byte(types.OpReject), nil)
}

func (c *Client) EndCall() error {
	c.mu.Lock()
	conn := c.conn
	c.activeCall = false
	c.callID = 0
	c.mediaKey = nil
	c.mu.Unlock()

	if conn == nil {
		return nil
	}

	st, err := conn.OpenStream()
	if err != nil {
		return err
	}
	defer st.Close()

	return protocol.WriteFrame(st, byte(types.OpEnd), nil)
}

func (c *Client) SendMedia(data []byte) error {
	c.mu.Lock()
	if !c.activeCall {
		c.mu.Unlock()
		return fmt.Errorf("no active call")
	}
	key := c.mediaKey
	cid := c.callID
	conn := c.conn
	c.mu.Unlock()

	if key == nil {
		return fmt.Errorf("no media key")
	}
	if conn == nil {
		return fmt.Errorf("not connected")
	}

	encrypted, err := encryptE2EE(key, data)
	if err != nil {
		return fmt.Errorf("encrypt: %w", err)
	}

	pkt := &protocol.MediaPacket{
		CallID:  cid,
		Payload: encrypted,
	}

	raw := protocol.MarshalMediaPacket(pkt)
	return conn.SendDatagram(raw)
}

func (c *Client) StartMediaLoop() {
	go c.handleDatagrams()
}

func (c *Client) deriveMediaKey() {
	shared, err := curve25519.X25519(c.e2eePriv[:], c.e2eePub[:])
	if err != nil {
		return
	}
	salt := make([]byte, 32)
	info := []byte("qcc-e2ee-media-key")
	kdf := hkdf.New(sha256.New, shared, salt, info)
	key := make([]byte, 32)
	if _, err := io.ReadFull(kdf, key); err != nil {
		return
	}
	c.mediaKey = key
}

type e2eeKeypair struct {
	PrivateKey [32]byte
	PublicKey  [32]byte
}

func generateE2EEKey() (*e2eeKeypair, error) {
	priv := make([]byte, 32)
	if _, err := io.ReadFull(rand.Reader, priv); err != nil {
		return nil, err
	}
	priv[0] &= 248
	priv[31] &= 127
	priv[31] |= 64

	pub, err := curve25519.X25519(priv, curve25519.Basepoint)
	if err != nil {
		return nil, err
	}

	var result e2eeKeypair
	copy(result.PrivateKey[:], priv)
	copy(result.PublicKey[:], pub)
	return &result, nil
}

func encryptE2EE(key, plaintext []byte) ([]byte, error) {
	aead, err := chacha20poly1305.New(key)
	if err != nil {
		return nil, err
	}
	nonce := make([]byte, chacha20poly1305.NonceSizeX)
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, err
	}
	return aead.Seal(nonce, nonce, plaintext, nil), nil
}

func decryptE2EE(key, ciphertext []byte) ([]byte, error) {
	aead, err := chacha20poly1305.New(key)
	if err != nil {
		return nil, err
	}
	if len(ciphertext) < chacha20poly1305.NonceSizeX {
		return nil, io.ErrUnexpectedEOF
	}
	nonce := ciphertext[:chacha20poly1305.NonceSizeX]
	return aead.Open(nil, nonce, ciphertext[chacha20poly1305.NonceSizeX:], nil)
}

func parseError(payload []byte) string {
	if len(payload) < 2 {
		return "unknown error"
	}
	return string(payload[2:])
}

func parseIdentityPayload(data []byte) (string, []byte, []byte) {
	if len(data) < 3 {
		return "", nil, nil
	}
	off := 0
	for off < len(data) && data[off] != 0 {
		off++
	}
	if off >= len(data) {
		return "", nil, nil
	}
	number := string(data[:off])
	off++

	if off+2 > len(data) {
		return number, nil, nil
	}
	certLen := int(binary.BigEndian.Uint16(data[off:]))
	off += 2
	if off+certLen > len(data) {
		return number, nil, nil
	}
	certDER := make([]byte, certLen)
	copy(certDER, data[off:off+certLen])
	off += certLen

	caCertDER := make([]byte, len(data)-off)
	copy(caCertDER, data[off:])

	return number, certDER, caCertDER
}

func parseRingPayload(data []byte) (string, [32]byte) {
	if len(data) < 1 {
		return "", [32]byte{}
	}
	numLen := int(data[0])
	if len(data) < 1+numLen+32 {
		return "", [32]byte{}
	}
	number := string(data[1 : 1+numLen])
	var e2eePub [32]byte
	copy(e2eePub[:], data[1+numLen:1+numLen+32])
	return number, e2eePub
}

func (c *Client) VerifyCertificate() bool {
	if len(c.caCertDER) == 0 || len(c.certDER) == 0 {
		return false
	}

	caCert, err := x509.ParseCertificate(c.caCertDER)
	if err != nil {
		return false
	}

	cert, err := x509.ParseCertificate(c.certDER)
	if err != nil {
		return false
	}

	roots := x509.NewCertPool()
	roots.AddCert(caCert)

	opts := x509.VerifyOptions{
		Roots: roots,
	}

	_, err = cert.Verify(opts)
	return err == nil
}