sock/server.go

package sock

import (
	"io"
	"net"
	"os"
	"path/filepath"

	"git.milar.in/milarin/cmap"
	"git.milar.in/milarin/slices"
)

type Server struct {
	socketPath string
	server     net.Listener
	clients    *cmap.Map[net.Conn, struct{}]

	onNewClient func(client net.Conn)
}

func Listen(socketPath string, onNewClient func(client net.Conn)) (*Server, error) {
	absPath, err := filepath.Abs(socketPath)
	if err != nil {
		return nil, err
	}

	os.Remove(absPath)
	server, err := net.Listen("unix", absPath)
	if err != nil {
		return nil, err
	}

	s := &Server{
		socketPath:  absPath,
		server:      server,
		clients:     cmap.New[net.Conn, struct{}](),
		onNewClient: onNewClient,
	}

	go s.handleClients()

	return s, nil
}

func (s *Server) Broadcast(r io.Reader) error {
	clients := slices.Map(s.clients.Keys(), func(c net.Conn) io.Writer { return c })
	w := newMultiWriter(clients...)

	if _, err := io.Copy(w, r); err != nil {
		return err
	}

	return nil
}

func (s *Server) handleClients() {
	for client, err := s.server.Accept(); err == nil; client, err = s.server.Accept() {
		go s.handleClient(client)
	}
}

func (s *Server) handleClient(client net.Conn) {
	s.clients.Put(client, struct{}{})
	defer s.clients.Delete(client)

	if s.onNewClient != nil {
		s.onNewClient(client)
	}

	data := make([]byte, 1024)
	for _, err := client.Read(data); err == nil; _, err = client.Read(data) {
	}
}

func (s *Server) Close() error {
	s.clients.Iter(func(client net.Conn, _ struct{}) { client.Close() })

	if err := s.server.Close(); err != nil {
		return err
	}

	if err := os.Remove(s.socketPath); err != nil {
		return err
	}

	return nil
}