go-chat/server/main.go

360 lines
9.5 KiB
Go
Raw Normal View History

2024-01-14 17:59:51 +00:00
package main
import (
"context"
"encoding/json"
"fmt"
"log"
"net"
"net/http"
"os"
"os/signal"
"path/filepath"
"strings"
"sync"
"time"
c "go-chat/common"
"github.com/alexflint/go-arg"
ws "nhooyr.io/websocket"
"nhooyr.io/websocket/wsjson"
)
type user struct {
room string
nick string
}
type conns struct {
sm sync.Mutex
cm map[*ws.Conn]user
}
type server struct {
admin string
logFn func(string, ...interface{})
conns *conns
rooms map[string]struct{}
rhist map[string][]c.SMsg
rhlen int
nickm map[string]string
}
func main() {
log := log.New(os.Stderr, "ws server 🚀 ", log.LstdFlags|log.Lshortfile|log.Lmsgprefix)
var args struct {
Admin string `arg:"-a" default:"8bit" help:"admin user nick, allows access to /sudo" placeholder:"NICK"`
2024-01-14 21:13:55 +00:00
HistLen uint `arg:"-l" default:"10" help:"set message history size" placeholder:"N"`
2024-01-14 17:59:51 +00:00
Port uint `arg:"positional" default:"0" help:"port to listen on, random available port if not set"`
NickMap *string `arg:"-n" help:"path to nick:pass JSON file" placeholder:"FILE"`
}
arg.MustParse(&args)
nickMap, err := loadNickMap(args.NickMap)
if err != nil {
log.Fatal(err)
}
err = run("localhost:"+fmt.Sprint(args.Port), nickMap, args.Admin, int(args.HistLen), log)
if err != nil {
log.Fatal(err)
}
}
func run(addr string, nickMap map[string]string, admin string, rhlen int, log *log.Logger) error {
listener, err := net.Listen("tcp4", addr)
if err != nil {
return err
}
log.Printf("listening on ws://%v", listener.Addr())
server := &http.Server{
Handler: server{
admin: admin,
logFn: log.Printf,
conns: &conns{cm: make(map[*ws.Conn]user)},
rooms: map[string]struct{}{"general": {}, "test1": {}, "test2": {}},
rhist: make(map[string][]c.SMsg),
rhlen: rhlen,
nickm: nickMap,
},
ReadTimeout: 10 * time.Second,
WriteTimeout: 10 * time.Second,
}
errch := make(chan error, 1)
go func() {
errch <- server.Serve(listener)
}()
signals := make(chan os.Signal, 1)
signal.Notify(signals, os.Interrupt)
select {
case err := <-errch:
log.Printf("failed to serve: %v", err)
case signal := <-signals:
log.Printf("quitting: %v", signal)
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
return server.Shutdown(ctx)
}
func (s server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
ctx := r.Context()
conn, err := ws.Accept(w, r, nil)
if err != nil {
s.logFn("%v", err)
return
}
defer conn.CloseNow()
if conn.Subprotocol() != "" {
return
}
port := strings.Split(r.RemoteAddr, ":")[1]
s.conns.sm.Lock()
s.conns.cm[conn] = user{room: "general", nick: port}
s.conns.sm.Unlock()
defer func() {
s.conns.sm.Lock()
delete(s.conns.cm, conn)
s.logFn("Remaining connections: %v", len(s.conns.cm))
s.conns.sm.Unlock()
}()
s.logFn("connected: %v", r.RemoteAddr)
for i := range s.rhist["general"] {
wsjson.Write(ctx, conn, s.rhist["general"][i])
}
cmsg := c.CMsg{}
smsg := c.SMsg{Id: port}
for {
err := func(ctx context.Context, conn *ws.Conn) error {
err := wsjson.Read(ctx, conn, &cmsg)
if err != nil {
return err
}
switch cmsg.Typ {
case c.Sudo:
s.logFn("(%v) sudo: %v", smsg.Id, cmsg.Msg)
if smsg.Id == s.admin {
cmd := strings.Split(cmsg.Msg, " ")
if len(cmd) == 2 {
if cmd[0] == "mk" {
if _, ok := s.rooms[cmd[1]]; ok {
wsjson.Write(ctx, conn, c.SMsg{Tim: time.Now(), Id: "system", Msg: fmt.Sprintf("Room exists: %v", cmd)})
} else {
s.rooms[cmd[1]] = struct{}{}
wsjson.Write(ctx, conn, c.SMsg{Tim: time.Now(), Id: "system", Msg: fmt.Sprintf("Created room: %v", cmd)})
}
} else if cmd[0] == "rm" {
if _, ok := s.rooms[cmd[1]]; ok && cmd[1] != "general" {
delete(s.rooms, cmd[1])
s.rhist[cmd[1]] = []c.SMsg{}
tim := time.Now()
s.conns.sm.Lock()
for cn, r := range s.conns.cm {
if r.room == cmd[1] {
r.room = "general"
s.conns.cm[cn] = r
wsjson.Write(ctx, cn, c.SMsg{Tim: tim, Id: "system", Msg: "room deleted, reconnected to general"})
}
}
s.conns.sm.Unlock()
wsjson.Write(ctx, conn, c.SMsg{Tim: time.Now(), Id: "system", Msg: fmt.Sprintf("Deleted room: %v", cmd)})
} else {
wsjson.Write(ctx, conn, c.SMsg{Tim: time.Now(), Id: "system", Msg: fmt.Sprintf("Room does not exist: %v", cmd)})
}
} else if cmd[0] == "yeet" {
found := false
s.conns.sm.Lock()
for c, r := range s.conns.cm {
if r.nick == cmd[1] {
c.Close(ws.StatusNormalClosure, "Kicked")
found = true
break
}
}
s.conns.sm.Unlock()
if found {
wsjson.Write(ctx, conn, c.SMsg{Tim: time.Now(), Id: "system", Msg: fmt.Sprintf("Yeet: %v", cmd[1])})
} else {
wsjson.Write(ctx, conn, c.SMsg{Tim: time.Now(), Id: "system", Msg: fmt.Sprintf("Not found: %v", cmd[1])})
}
} else {
wsjson.Write(ctx, conn, c.SMsg{Tim: time.Now(), Id: "system", Msg: fmt.Sprintf("Invalid command: %v", cmd)})
}
} else if cmd[0] == "wc" {
s.conns.sm.Lock()
wc := len(s.conns.cm)
s.conns.sm.Unlock()
wsjson.Write(ctx, conn, c.SMsg{Tim: time.Now(), Id: "system", Msg: fmt.Sprintf("Online: %v", wc)})
} else if cmd[0] == "man" {
wsjson.Write(ctx, conn, c.SMsg{Tim: time.Now(), Id: "system", Msg: "Available commands: man, mk, rm, wc, yeet"})
} else {
wsjson.Write(ctx, conn, c.SMsg{Tim: time.Now(), Id: "system", Msg: fmt.Sprintf("Invalid command: %v", cmd)})
}
} else {
wsjson.Write(ctx, conn, c.SMsg{Tim: time.Now(), Id: "system", Msg: "Unrecognised command, use /man for more info"})
}
case c.Echo:
s.logFn("(%v) echo: %v", smsg.Id, cmsg.Msg)
s.conns.sm.Lock()
room := s.conns.cm[conn].room
s.conns.sm.Unlock()
smsg.Tim = time.Now()
smsg.Msg = cmsg.Msg
if len(s.rhist[room]) < s.rhlen {
s.rhist[room] = append(s.rhist[room], smsg)
} else {
s.rhist[room] = append(s.rhist[room][1:], smsg)
}
s.conns.sm.Lock()
for c, r := range s.conns.cm {
if r.room == room {
wsjson.Write(ctx, c, &smsg)
}
}
s.conns.sm.Unlock()
case c.Mv:
switch nick, valid := verifyNick(&s, cmsg.Msg); valid {
case nickOk:
s.logFn("(%v) mv: %v", smsg.Id, cmsg.Msg)
smsg.Id = nick
s.conns.sm.Lock()
u := s.conns.cm[conn]
u.nick = smsg.Id
s.conns.cm[conn] = u
s.conns.sm.Unlock()
wsjson.Write(ctx, conn, c.SMsg{Tim: time.Now(), Id: "system", Msg: fmt.Sprintf("nick set: %v", nick)})
case nickUsed:
s.logFn("(%v) mv used: %v", smsg.Id, cmsg.Msg)
wsjson.Write(ctx, conn, c.SMsg{Tim: time.Now(), Id: "system", Msg: fmt.Sprintf("nick in use: %v", cmsg.Msg)})
case nickInvalid:
s.logFn("(%v) mv invalid: %v", smsg.Id, cmsg.Msg)
wsjson.Write(ctx, conn, c.SMsg{Tim: time.Now(), Id: "system", Msg: fmt.Sprintf("invalid nick: %v", cmsg.Msg)})
}
case c.Ls:
s.logFn("(%v) ls", smsg.Id)
s.conns.sm.Lock()
room := s.conns.cm[conn].room
s.conns.sm.Unlock()
avRooms := ""
for r := range s.rooms {
avRooms += r + ", "
}
wsjson.Write(ctx, conn, c.SMsg{Tim: time.Now(), Id: "system", Msg: fmt.Sprintf("connected to: %v, available: %v", room, avRooms[:len(avRooms)-2])})
case c.Cd:
if _, ok := s.rooms[cmsg.Msg]; ok {
s.logFn("(%v) cd: %v", smsg.Id, cmsg.Msg)
s.conns.sm.Lock()
u := s.conns.cm[conn]
u.room = cmsg.Msg
s.conns.cm[conn] = u
s.conns.sm.Unlock()
wsjson.Write(ctx, conn, c.SMsg{Tim: time.Now(), Id: "system", Msg: fmt.Sprintf("connected to: %v", u.room)})
rhistSize := max(len(s.rhist[u.room])-10, 0)
recentHistory := s.rhist[u.room][rhistSize:]
for i := range recentHistory {
wsjson.Write(ctx, conn, recentHistory[i])
}
} else {
s.logFn("(%v) cd invalid: %v", smsg.Id, cmsg.Msg)
wsjson.Write(ctx, conn, c.SMsg{Tim: time.Now(), Id: "system", Msg: fmt.Sprintf("unchanged, invalid room: %v", cmsg.Msg)})
}
case c.Who:
s.conns.sm.Lock()
room := s.conns.cm[conn].room
s.logFn("(%v) who: %v", smsg.Id, room)
users := fmt.Sprintf("users in %v: ", room)
for _, r := range s.conns.cm {
if r.room == room {
users += fmt.Sprintf("%v, ", r.nick)
}
}
s.conns.sm.Unlock()
wsjson.Write(ctx, conn, &c.SMsg{Tim: time.Now(), Id: "system", Msg: users[:len(users)-2]})
}
return nil
}(ctx, conn)
if ws.CloseStatus(err) == ws.StatusNormalClosure {
s.logFn("disconnected: %v", r.RemoteAddr)
return
}
if err != nil {
s.logFn("failed, addr %v: %v", r.RemoteAddr, err)
return
}
}
}
type nickErr int
const (
nickOk nickErr = iota
nickUsed
nickInvalid
)
func verifyNick(s *server, n string) (string, nickErr) {
nick, pass, _ := strings.Cut(n, ":")
s.conns.sm.Lock()
defer s.conns.sm.Unlock()
for _, u := range s.conns.cm {
if u.nick == nick {
return "", nickUsed
}
}
expPass, needAuth := s.nickm[nick]
if (!needAuth || pass == expPass) && alphanumeric(nick) {
return nick, nickOk
} else {
return "", nickInvalid
}
}
func loadNickMap(m *string) (map[string]string, error) {
nm := make(map[string]string)
if m == nil {
return nm, nil
}
path, err := filepath.Abs(*m)
if err != nil {
return nil, err
}
file, err := os.ReadFile(path)
if err != nil {
return nil, err
}
err = json.Unmarshal(file, &nm)
if err != nil {
return nil, err
}
return nm, nil
}
func alphanumeric(s string) bool {
for _, r := range s {
if (r < 'a' || r > 'z') && (r < 'A' || r > 'Z') && (r < '0' || r > '9') {
return false
}
}
return true
}