Last active
November 5, 2023 12:30
-
-
Save YabZhang/97594b471926d61b567b7ba23c98fadf to your computer and use it in GitHub Desktop.
smallchat.go
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| package main | |
| import ( | |
| "bytes" | |
| "flag" | |
| "fmt" | |
| "net" | |
| "sync" | |
| ) | |
| const MAX_CLIENTS = 1000 | |
| var serverPort int | |
| var chatState *ChatState | |
| type ChatState struct { | |
| mux sync.RWMutex | |
| num int | |
| clients map[*Client]struct{} | |
| } | |
| func (cs *ChatState) addClient(client *Client) { | |
| cs.mux.Lock() | |
| defer cs.mux.Unlock() | |
| cs.num += 1 | |
| cs.clients[client] = struct{}{} | |
| } | |
| func (cs *ChatState) removeClient(client *Client) { | |
| cs.mux.Lock() | |
| defer cs.mux.Unlock() | |
| cs.num -= 1 | |
| delete(cs.clients, client) | |
| } | |
| func (cs *ChatState) broadcast(mesg []byte, exclude *Client) { | |
| cs.mux.RLock() | |
| defer cs.mux.RUnlock() | |
| for client, _ := range cs.clients { | |
| if client == nil || client == exclude { | |
| continue | |
| } | |
| _, err := client.Conn.Write(mesg) | |
| if err != nil { | |
| fmt.Printf("Error writing to conn: %+v\n", err) | |
| } | |
| } | |
| } | |
| type Client struct { | |
| Nick string | |
| Conn net.Conn | |
| } | |
| func init() { | |
| flag.IntVar(&serverPort, "port", 7789, "Port to listen on") | |
| chatState = &ChatState{ | |
| clients: make(map[*Client]struct{}, MAX_CLIENTS), | |
| } | |
| } | |
| func handleNewConn(client *Client) { | |
| client.Conn.Write([]byte("Welcome to Simple Chat! Use /nick to change nick name.\n")) | |
| addr := client.Conn.RemoteAddr().String() | |
| fmt.Printf("New connection from %s\n", addr) | |
| client.Nick = "user: " + addr | |
| for { | |
| buf := make([]byte, 256) | |
| nread, err := client.Conn.Read(buf) | |
| if err != nil { | |
| fmt.Printf("Error reading from conn: %+v\n", err) | |
| chatState.removeClient(client) | |
| return | |
| } | |
| mesg := bytes.TrimSpace(buf[:nread]) | |
| fmt.Printf("%s: %s\n", client.Nick, mesg) | |
| if len(mesg) <= 0 { | |
| continue | |
| } | |
| if mesg[0] == '/' { | |
| parts := bytes.SplitN(mesg, []byte(" "), 2) | |
| switch string(parts[0]) { | |
| case "/nick": | |
| fmt.Printf("content: %#v\n", parts) | |
| if len(parts) != 2 { | |
| fmt.Printf("Invalid command: %#v\n", mesg) | |
| continue | |
| } | |
| client.Nick = string(parts[1]) | |
| fmt.Printf("result: %#v\n", client.Nick) | |
| } | |
| continue | |
| } | |
| chatState.broadcast([]byte(client.Nick+"> "+string(mesg)+"\n"), client) | |
| } | |
| } | |
| func main() { | |
| flag.Parse() | |
| fmt.Println("Starting server on port", serverPort) | |
| ln, err := net.Listen("tcp", fmt.Sprintf("0.0.0.0:%d", serverPort)) | |
| if err != nil { | |
| panic(err) | |
| } | |
| for { | |
| conn, err := ln.Accept() | |
| if err != nil { | |
| fmt.Printf("Error accepting conn: %+v\n", err) | |
| continue | |
| } | |
| client := &Client{Conn: conn} | |
| chatState.addClient(client) | |
| go handleNewConn(client) | |
| } | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment