package main import ( "errors" "fmt" "net/http" "os" "syscall" "time" "github.com/gorilla/websocket" "github.com/joho/godotenv" "github.com/labstack/echo/v4" "github.com/labstack/echo/v4/middleware" "github.com/labstack/gommon/log" ) const TimeFormat = "2006-01-02 15:04:05" type Client struct { ID int conn *websocket.Conn Exit bool } var id = 0 var clients = []*Client{} var keyDir = "." var token = "" var upgrader = websocket.Upgrader{} type ReqData struct { ID int `json:"id"` } func main() { _ = godotenv.Load() keyDir = os.Getenv("KEY_DIR") if keyDir == "" { log.Fatal("KEY_DIR is not defined") } token = os.Getenv("TOKEN") if token == "" { log.Fatal("TOKEN is not defined") } app := echo.New() app.Logger.SetLevel(log.INFO) app.Use(middleware.RecoverWithConfig(middleware.RecoverConfig{ StackSize: 1 << 10, LogLevel: log.ERROR, })) app.Use(middleware.Secure()) app.Use(middleware.LoggerWithConfig(middleware.LoggerConfig{ Format: "${time_custom} ${method} ${uri} ---> ${status} in ${latency_human} (${bytes_out} bytes)", CustomTimeFormat: TimeFormat, })) app.Use(middleware.RemoveTrailingSlash()) app.OnAddRouteHandler = func(host string, route echo.Route, handler echo.HandlerFunc, middleware []echo.MiddlewareFunc) { now := time.Now() fmt.Printf("%s registered %-6s %s\n", now.Format(TimeFormat), route.Method, route.Path) } app.HTTPErrorHandler = func(err error, c echo.Context) { if c.Response().Committed { log.Error(err) return } code := http.StatusInternalServerError var response any = err.Error() var httpErr *echo.HTTPError if errors.As(err, &httpErr) { code = httpErr.Code response = httpErr } if code >= 500 { log.Error(err) } err2 := c.JSON(code, response) if err2 != nil { log.Error(err2) } } app.GET("/", func(c echo.Context) error { return c.JSON(http.StatusOK, "Hello World!") }) app.GET("/keys", keys) app.GET("/clients", func(c echo.Context) error { if c.Request().Header.Get("Authorization") != token { return c.NoContent(http.StatusUnauthorized) } var a = []echo.Map{} for _, client := range clients { a = append(a, echo.Map{ "id": client.ID, "exit": client.Exit, }) } return c.JSON(http.StatusOK, a) }) app.POST("/exit", func(c echo.Context) error { if c.Request().Header.Get("Authorization") != token { return c.NoContent(http.StatusUnauthorized) } var data ReqData err := c.Bind(&data) if err != nil { return c.JSON(http.StatusBadRequest, echo.Map{"message": "missing field 'id'"}) } for _, client := range clients { if client.ID == data.ID { client.Exit = true return c.JSON(http.StatusOK, echo.Map{"message": "ok"}) } } return c.JSON(http.StatusNotFound, echo.Map{"message": "not found"}) }) log.Fatal(app.Start(":8080")) } func keys(c echo.Context) error { ws, err := upgrader.Upgrade(c.Response(), c.Request(), nil) if err != nil { return err } defer ws.Close() client := Client{ ID: id, conn: ws, Exit: false, } clients = append(clients, &client) id += 1 fmt.Printf("%s client %-3d connected", time.Now().Format(TimeFormat), client.ID) err = ws.WriteMessage(websocket.TextMessage, []byte("welcome")) if err != nil { log.Error(err) } f, err := os.OpenFile(fmt.Sprintf("%s/%s_%d.txt", keyDir, time.Now().Format("2006-01-02_15-04"), client.ID), syscall.O_CREAT|syscall.O_APPEND|syscall.O_WRONLY, 0644) if err != nil { log.Error(err) fmt.Printf("%s client %-3d crashed (couldn't open file)", time.Now().Format(TimeFormat), client.ID) client.Exit = true return c.NoContent(http.StatusInternalServerError) } defer f.Close() for { if client.Exit { _ = ws.WriteMessage(websocket.TextMessage, []byte("exit")) _ = ws.Close() fmt.Printf("%s client %-3d disconnected", time.Now().Format(TimeFormat), client.ID) client.Exit = true break } _, msg, err := ws.ReadMessage() if err != nil { log.Error(err) _ = ws.WriteMessage(websocket.TextMessage, []byte("exit")) _ = ws.Close() fmt.Printf("%s client %-3d crashed (websocket error)", time.Now().Format(TimeFormat), client.ID) client.Exit = true break } _, err = f.Write(msg) if err != nil { log.Error(err) _ = ws.WriteMessage(websocket.TextMessage, []byte("exit")) _ = ws.Close() fmt.Printf("%s client %-3d crashed (file write error)", time.Now().Format(TimeFormat), client.ID) client.Exit = true break } } return c.NoContent(http.StatusOK) }