go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/experiments/megaphone/main.go

/*

Copyright (c) 2023 - Present. Will Charczuk. All rights reserved.
Use of this source code is governed by a MIT license that can be found in the LICENSE file at the root of the repository.

*/

package main

import (
	"bytes"
	"context"
	"encoding/gob"
	"fmt"
	"net"
	"os"
	"os/user"
	"time"

	"github.com/urfave/cli/v2"
	"go.charczuk.com/sdk/collections"
)

var (
	ipv4mcastaddr = &net.UDPAddr{
		IP:   net.ParseIP("224.0.0.251"),
		Port: 1337,
	}
)

const bufferSize = 32 * 1024

func main() {
	app := &cli.App{
		Name:      "mega",
		Usage:     "Megaphone is a UDP anycast based chat program.",
		Copyright: "(c) 2023 Will Charczuk",
		Commands: []*cli.Command{
			listen(),
			send(),
		},
	}
	if err := app.Run(os.Args); err != nil {
		fmt.Fprintf(os.Stderr, "%+v\n", err)
		os.Exit(1)
	}
}

func listen() *cli.Command {
	return &cli.Command{
		Name:  "listen",
		Usage: "Print message received",
		Action: func(ctx *cli.Context) error {
			c := new(Client)
			if err := c.Open(); err != nil {
				return err
			}

			c.Listen(ctx.Context)
			var msg Message
			for {
				msg = <-c.Messages
				fmt.Printf("%s\t%s\t%s\n", msg.Timestamp.Format(time.Kitchen), msg.User, msg.Message)
			}
		},
	}
}

func send() *cli.Command {
	return &cli.Command{
		Name:  "send",
		Usage: "Send a message",
		Action: func(ctx *cli.Context) error {
			user, err := user.Current()
			if err != nil {
				return err
			}
			host, err := os.Hostname()
			if err != nil {
				return err
			}
			message := Message{
				User:      fmt.Sprintf("%s@%s", user.Username, host),
				Timestamp: time.Now().UTC(),
				Message:   ctx.Args().First(),
			}
			c := new(Client)
			if err = c.Open(); err != nil {
				return err
			}
			c.SendMessage(ctx.Context, message)

			fmt.Println("sent!")
			return nil
		},
	}
}

type Client struct {
	Messages chan Message

	recv     *net.UDPConn
	send     *net.UDPConn
	peers    map[string]*PeerInfo
	messages *collections.Queue[Message]
}

func (c *Client) Open() error {
	c.messages = new(collections.Queue[Message])
	var err error
	c.recv, err = net.ListenMulticastUDP("udp4", nil, ipv4mcastaddr)
	if err != nil {
		return fmt.Errorf("client cannot open receive connection; %w", err)
	}
	c.recv.SetReadBuffer(bufferSize)
	c.send, err = net.DialUDP("udp4", nil, ipv4mcastaddr)
	if err != nil {
		return err
	}
	c.send.SetWriteBuffer(bufferSize)

	// do initial sync
	return c.startSync()
}

func (c *Client) Listen(ctx context.Context) {
	c.Messages = make(chan Message)
	go func() {
		var readBytes int
		var packet Packet
		for {
			select {
			case <-ctx.Done():
				return
			default:
			}
			buffer := make([]byte, bufferSize)
			readBytes, _, _ = c.recv.ReadFromUDP(buffer)
			_ = gob.NewDecoder(bytes.NewBuffer(buffer[:readBytes])).Decode(&packet)
			c.handleIncomingPacket(packet)
		}
	}()
}

func (c *Client) SendMessage(ctx context.Context, msg Message) error {
	return nil
}

func (c *Client) Close() error {
	if c.recv != nil {
		_ = c.Close()
	}
	if c.send != nil {
		_ = c.send.Close()
	}
	return nil
}

//
// internal client methods
//

func (c *Client) startSync() error {
	return nil
}

func (c *Client) handleIncomingPacket(pck Packet) {
	switch pck.Method {
	case MethodHeartbeat:
		heartbeat, _ := pck.Heartbeat()
		if c.peers == nil {
			c.peers = make(map[string]*PeerInfo)
		}
		c.peers[heartbeat.User].LastSeen = time.Now().UTC()
	case MethodSync:
		// nop
	}
}

func (c *Client) sendPacket(ctx context.Context, pck Packet) error {
	buf := new(bytes.Buffer)
	err := gob.NewEncoder(buf).Encode(pck)
	if err != nil {
		return fmt.Errorf("client cannot send packet; %w", err)
	}
	if _, err = c.send.Write(buf.Bytes()); err != nil {
		return err
	}
	return nil
}

type Packet struct {
	Version uint8
	Size    uint16
	Method  uint16
	Data    []byte
}

// PeerInfo is metadata about peers.
type PeerInfo struct {
	User     string
	LastSeen time.Time
}

// Message is a byte packed message we'll send or receive over the wire.
type Message struct {
	User      string
	Hostname  string
	Timestamp time.Time
	Message   string
	Labels    map[string]string
}

type Heartbeat struct {
	User      string
	Hostname  string
	Timestamp time.Time
}

type SyncRequest struct {
	User string
	AsOf time.Time
}

func (p Packet) Message() (msg Message, err error) {
	err = gob.NewDecoder(bytes.NewBuffer(p.Data)).Decode(&msg)
	return
}

func (p Packet) Heartbeat() (hb Heartbeat, err error) {
	err = gob.NewDecoder(bytes.NewBuffer(p.Data)).Decode(&hb)
	return
}

type Method uint16

const (
	MethodSendMessage = 10
	MethodSync        = 20
	MethodHeartbeat   = 50
)