converge/cmd/converge/server.go

package main

import (
	"cidebug/pkg/iowrappers"
	"fmt"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"regexp"
	"sync"
)

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
}

func handleWebSocket(w http.ResponseWriter, r *http.Request,
	handler func(w http.ResponseWriter, r *http.Request, websockerConnection *iowrappers.WebSocketConn)) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println("Error upgrading to WebSocket:", err)
		return
	}
	wsConn := iowrappers.NewWebSocketConn(conn)
	defer wsConn.Close()

	handler(w, r, wsConn)
}

type WebSocketService struct {
	handler func(w http.ResponseWriter, r *http.Request, conn *iowrappers.WebSocketConn)
}

func (endpoint *WebSocketService) handle(w http.ResponseWriter, r *http.Request) {
	handleWebSocket(w, r, endpoint.handler)
}

type Agent struct {
	agent          *iowrappers.WebSocketConn
	publicId       string
	client         *iowrappers.WebSocketConn
	agentAvailable chan bool
	clientClosed   chan bool
}

func NewAgent(publicId string,
	agentConn *iowrappers.WebSocketConn,
	clientConn *iowrappers.WebSocketConn) *Agent {
	return &Agent{
		agent:          agentConn,
		publicId:       publicId,
		client:         clientConn,
		agentAvailable: make(chan bool, 1),
		clientClosed:   make(chan bool, 1),
	}
}

type Admin struct {
	// map of public id to agent
	mutex  sync.Mutex
	agents map[string]*Agent
}

func NewAdmin() *Admin {
	admin := Admin{
		mutex:  sync.Mutex{},
		agents: make(map[string]*Agent),
	}
	return &admin
}

func (admin *Admin) addAgent(publicId string, conn *iowrappers.WebSocketConn) (*Agent, error) {
	admin.mutex.Lock()
	defer admin.mutex.Unlock()

	agent := admin.agents[publicId]
	if agent != nil && agent.agent != nil && agent.agent != conn {
		return nil, fmt.Errorf("A different agent with same publicId '%s' already registered", publicId)
	}
	if agent != nil {
		agent.agent = conn
		return agent, nil
	}
	agent = NewAgent(publicId, conn, nil)
	admin.agents[publicId] = agent
	return agent, nil
}

func (admin *Admin) addClient(publicId string, conn *iowrappers.WebSocketConn) (*Agent, error) {
	admin.mutex.Lock()
	defer admin.mutex.Unlock()

	agent := admin.agents[publicId]
	if agent != nil && agent.client != nil && agent.client != conn {
		// we should setup on-demend connections ot agents later.
		return nil, fmt.Errorf("A different client with same publicId '%s' already registered", publicId)
	}
	if agent != nil {
		agent.client = conn
		return agent, nil
	}
	agent = NewAgent(publicId, nil, conn)
	admin.agents[publicId] = agent
	return agent, nil
}

func (admin *Admin) RemoveAgent(publicId string) error {
	admin.mutex.Lock()
	defer admin.mutex.Unlock()

	agent := admin.agents[publicId]
	if agent == nil {
		return fmt.Errorf("Cannot remove agent: '%s' not found", publicId)
	}
	agent.agent = nil
	if agent.client == nil {
		log.Printf("Removing agent: '%s'", publicId)
		delete(admin.agents, publicId)
	}
	return nil
}

func (admin *Admin) RemoveClient(publicId string) error {
	admin.mutex.Lock()
	defer admin.mutex.Unlock()

	agent := admin.agents[publicId]
	if agent == nil {
		return fmt.Errorf("Cannot remove agent: '%s' not found", publicId)
	}
	agent.client = nil
	if agent.agent == nil {
		log.Printf("Removing client: '%s'", publicId)
		delete(admin.agents, publicId)
	}
	return nil
}

func (admin *Admin) Register(publicId string, conn *iowrappers.WebSocketConn) error {
	defer conn.Close()
	agent, err := admin.addAgent(publicId, conn)
	if err != nil {
		return err
	}
	defer func() {
		admin.RemoveAgent(publicId)
	}()
	log.Printf("After defer remove agent\n")
	agent.agentAvailable <- true
	log.Printf("Agent registered: '%s'\n", publicId)
	<-agent.clientClosed
	return nil
}

func (admin *Admin) Connect(publicId string, conn *iowrappers.WebSocketConn) error {
	defer conn.Close()
	agent, err := admin.addClient(publicId, conn)
	if err != nil {
		return err
	}
	defer func() {
		admin.RemoveClient(publicId)
	}()
	<-agent.agentAvailable
	log.Printf("Connecting client and agent: '%s'\n", publicId)
	iowrappers.SynchronizeStreams(agent.client, agent.agent)
	agent.clientClosed <- true
	return nil
}

func (admin *Admin) log() {
	log.Println("CONNECTIONS")
	for _, agent := range admin.agents {
		log.Println(agent.publicId)
	}
	log.Printf("\n")
}

func parsePublicId(path string) (publicId string, _ error) {
	pattern := regexp.MustCompile("^/[^/]+/([^/]+)$")
	matches := pattern.FindStringSubmatch(path)
	if len(matches) != 2 {
		return "", fmt.Errorf("Invalid URL path '%s'", path)
	}
	return matches[1], nil

}

func main() {

	admin := NewAdmin()
	registrationService := WebSocketService{
		handler: func(w http.ResponseWriter, r *http.Request, conn *iowrappers.WebSocketConn) {
			publicId, err := parsePublicId(r.URL.Path)
			if err != nil {
				log.Printf("Cannot parse public id from url: '%v'\n", err)
				return
			}
			log.Printf("Got registration connection: '%s'\n", publicId)
			err = admin.Register(publicId, conn)
			if err != nil {
				log.Printf("Error %v\n", err)
			}
		},
	}
	clientService := WebSocketService{
		handler: func(w http.ResponseWriter, r *http.Request, conn *iowrappers.WebSocketConn) {
			publicId, err := parsePublicId(r.URL.Path)
			if err != nil {
				log.Printf("Cannot parse public id from url: '%v'\n", err)
				return
			}
			log.Printf("Got client connection: '%s'\n", publicId)
			err = admin.Connect(publicId, conn)
			if err != nil {
				log.Printf("Error %v\n", err)
			}
		},
	}

	http.HandleFunc("/register/", registrationService.handle)
	http.HandleFunc("/client/", clientService.handle)

	// Start HTTP server
	fmt.Println("WebSocket server listening on :8000")
	log.Fatal(http.ListenAndServe(":8000", nil))
}