server.go

package main

import (
	"bufio"
	"fmt"
	"math/rand"
	"net"
	"strings"
	"time"
)

func writeData(c net.Conn, data []byte) {
	_, err := c.Write(data)
	if err != nil {
		fmt.Println(err)
	}
}

// https://flaviocopes.com/go-random/
func randomInt(min, max int) int {
	return min + rand.Intn(max-min)
}

func randomString(len int) string {
	bytes := make([]byte, len)
	for i := 0; i < len; i++ {
		bytes[i] = byte(randomInt(65, 90))
	}
	return string(bytes)
}

type User struct {
	id        string
	conn      net.Conn
	publicKey string
	receiver  string
}

func readData(c net.Conn) string {
	netData, err := bufio.NewReader(c).ReadString('\n')
	if err != nil {
		fmt.Println(err)
		return ""
	}

	m := strings.TrimSpace(string(netData))
	fmt.Printf("\n" + m + "\n")
	return m
}
func initConn(userID string, users map[string]*User) {
	c := users[userID].conn

	m := readData(c)
	if m == "pk" {
		writeData(c, []byte("ok\n"))
		pk := readData(c)
		users[userID].publicKey = pk
		s := readData(c)
		writeData(c, []byte("ok\n"))
		if s == "wu" {
			return
		}
		writeData(c, []byte("to\n"))
		to := readData(c)
		if _, ok := users[m]; ok && len(users[m].receiver) == 0 {
			users[userID].receiver = to
			users[to].receiver = userID
			users[to].conn.Write([]byte("uc\n"))
			// send public keys
			writeData(users[userID].conn, []byte(users[to].publicKey+"\n"))
			writeData(users[to].conn, []byte(users[userID].publicKey+"\n"))
		}
	} else {
		c.Close()
		delete(users, userID)
	}
}

func handleMess(userID string, mess string, users map[string]*User) {
	if to, ok := users[users[userID].receiver]; ok {
		to.conn.Write([]byte(mess + "\n"))
	} else {
		users[userID].conn.Write([]byte("$Fatal error: The connection to the recipient could not be found, so the message was not delivered to him." + "\n"))
		users[userID].conn.Close()
		delete(users, userID)
	}
}

func handleConn(c net.Conn, users map[string]*User) {
	u := User{id: randomString(8), conn: c}

	users[u.id] = &u

	hm := readData(c)
	if hm == "hi server" {
		writeData(c, []byte("hi client"+"\n"))
	} else {
		return
	}
	defer c.Close()
	defer func(userID string, users map[string]*User) {
		if _, ok := users[userID]; ok {
			if len(users[userID].receiver) != 0 {
				receiver := users[userID].receiver
				users[receiver].conn.Close()
				delete(users, receiver)
			}
			delete(users, userID)
		}

	}(u.id, users)
	// send ID
	writeData(c, []byte(u.id+"\n"))

	initConn(u.id, users)
	for {
		netData, err := bufio.NewReader(c).ReadString('\n')
		if err != nil {
			fmt.Println(err)
			return
		}

		temp := strings.TrimSpace(string(netData))
		if temp == "STOP" {
			break
		}
		handleMess(u.id, temp, users)
	}

}

// https://opensource.com/article/18/5/building-concurrent-tcp-server-go
func main() {
	users := make(map[string]*User)
	rand.Seed(time.Now().UnixNano())
	l, err := net.Listen("tcp4", ":8221")
	if err != nil {
		fmt.Println(err)
		return
	}
	defer l.Close()

	for {
		c, err := l.Accept()
		if err != nil {
			fmt.Println(err)
			break
		}
		go handleConn(c, users)
	}
}