chitter.go

package main

import (
	"bufio"
	"fmt"
	"net"
	"os"
	"strconv"
	"strings"
)

var idAssignmentChan = make(chan string)
var registerChan = make(chan Client)
var clients = make(map[string]net.Conn)
var broadcast = make(chan string)
var private = make(chan string)

type Client struct {
	con net.Conn
	id  string
}

func registerConnection() {
	for {
		select {
		case c := <-registerChan:
			clients[c.id] = c.con

		case m := <- broadcast:
			for _, connection := range clients {
				connection.Write([]byte(m))
			}

		case m := <- private:
			id, content := splitStrings(m)
			if _, ok := clients[id]; ok {
				clients[id].Write([]byte(content))
			}
		}
	}
}

func splitStrings(line string) (string, string) {
	splits := strings.SplitAfterN(string(line)+": ", ":", 2)
	tmp0 := splits[0]
	tmp1 := splits[1]

	id := tmp0[:len(tmp0)-1]

	content := tmp1
	if len(tmp1) > 1 {
		content = tmp1[:len(tmp1)-2]
	}
	return id, content
}

func HandleConnection(conn net.Conn) {
	b := bufio.NewReader(conn)
	client_id := <-idAssignmentChan
	registerChan <- Client{conn, client_id}

	for {
		line, err := b.ReadBytes('\n')
		if err != nil {
			conn.Close()
			break
		}

		id, content := splitStrings(string(line))

		switch id {
		case "whoami":
			conn.Write([]byte("chitter: " + client_id + "\n"))
		case "all":
			broadcast <- string(client_id + ": " + content)
		default:
			private <- string(id + ":" + client_id + ":" + content)
		}
	}
}

func IdManager() {
	var i uint64
	for i = 0; ; i++ {
		idAssignmentChan <- strconv.FormatUint(i, 10)
	}
}

func main() {
	if len(os.Args) < 2 {
		fmt.Fprintf(os.Stderr, "Usage: chitter <port-number>\n")
		os.Exit(1)
		return
	}
	port := os.Args[1]
	server, err := net.Listen("tcp", ":"+port)
	if err != nil {
		fmt.Fprintln(os.Stderr, "Can't connect to port")
		os.Exit(1)
	}

	go IdManager()
	go registerConnection()
	fmt.Println("Listening on port", os.Args[1])
	for {
		conn, _ := server.Accept()
		go HandleConnection(conn)
	}
}