TCP line-based middleware in Go

middleware.go

package main
/**
 * Simple line-based TCP middleware.
 *
 * Server protocol:
 * 		send: "server"
 *		send: server_id
 *	LOOP:
 * 		recv: request
 *		send: response
 *
 * Client protocol:
 *		send: "client"
 *	LOOP:
 *		send: server_id
 *		send: request
 *		recv: response
 */

import (
	"bufio"
	"fmt"
	"log"
	"net"
	"sync"
)

type Request struct {
	Data  string
	Reply chan string
}

type Server struct {
	Req     chan Request
	Stop    chan int
	Stopped chan int
}

var ServerMap sync.Map // map[string]Server

func main() {
	// Start TCP server
	l, err := net.Listen("tcp", ":8080")
	if err != nil {
		log.Fatal("listen:", err)
	}
	defer l.Close()

	log.Println("Serving middleware on :8080")
	for {
		// Accept incoming connections
		conn, err := l.Accept()
		if err != nil {
			log.Fatal("accept:", err)
		}
		go handle(conn)
	}
}

func handle(conn net.Conn) {
	defer conn.Close()

	// Read first line
	scanner := bufio.NewScanner(conn)
	if !scanner.Scan() {
		fmt.Fprintf(conn, "error\n")
		return
	}

	// Dispatch to appropriate handler
	switch scanner.Text() {
	case "server":
		handleServer(conn, scanner)
	case "client":
		handleClient(conn, scanner)
	default:
		fmt.Fprintf(conn, "error\n")
		return
	}
}

func handleServer(conn net.Conn, scanner *bufio.Scanner) {
	// First line is ID
	if !scanner.Scan() {
		fmt.Fprintf(conn, "error\n")
		return
	}
	id := scanner.Text()

	// Save server to global map
	server := Server{
		Req:     make(chan Request),
		Stop:    make(chan int, 1),
		Stopped: make(chan int),
	}
	previous, loaded := ServerMap.Swap(id, server)

	// If server already connected, close it
	if loaded {
		select {
		case previous.(Server).Stop <- 0:
		default:
		}
	}

	// Cleanup server on exit
	defer func() {
		ServerMap.CompareAndDelete(id, server)
		close(server.Stopped)
		conn.Close()
	}()

	// Receive requests and send responses
	for {
		select {
		case req := <-server.Req:
			// Send received request
			_, err := fmt.Fprintf(conn, "%s\n", req.Data)
			// Exit on send error
			if err != nil {
				return
			}
			// Read response
			if !scanner.Scan() {
				req.Reply <- "error\n"
			} else {
				req.Reply <- scanner.Text()
			}
		case <-server.Stop:
			return
		}
	}
}

func handleClient(conn net.Conn, scanner *bufio.Scanner) {
	for {
		// First line is server Id
		if !scanner.Scan() {
			return
		}
		id := scanner.Text()

		// Second line is request
		if !scanner.Scan() {
			return
		}
		data := scanner.Text()

		// Get server from map
		value, ok := ServerMap.Load(id)
		if !ok {
			_, err := fmt.Fprintf(conn, "error\n")
			if err != nil {
				return
			}
			continue
		}
		server := value.(Server)

		// Try to send request
		req := Request{
			Data:  data,
			Reply: make(chan string, 1),
		}
		select {
		case server.Req <- req:
		case <-server.Stopped:
			_, err := fmt.Fprintf(conn, "error\n")
			if err != nil {
				return
			}
			continue
		}
		// Try to receive response
		select {
		case resp := <-req.Reply:
			_, err := fmt.Fprintf(conn, "%s\n", resp)
			if err != nil {
				return
			}
		case <-server.Stopped:
			_, err := fmt.Fprintf(conn, "error\n")
			if err != nil {
				return
			}
		}
	}
}