keys/server/main.go

package main

import (
	"errors"
	"fmt"
	"net/http"
	"os"
	"path"
	"syscall"
	"time"

	"github.com/google/uuid"
	"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 {
	Obj        *Dummy          `json:"obj"`
	conn       *websocket.Conn `json:"-"`
	ExitWanted bool            `json:"exitWanted"`
}

var clients = []*Client{}
var dataDir = "."
var token = ""
var upgrader = websocket.Upgrader{}

type ReqData struct {
	ID uuid.UUID `json:"id"`
}

type IDParam struct {
	ID string `param:"id"`
}

func authed(next echo.HandlerFunc) echo.HandlerFunc {
	return func(c echo.Context) error {
		if c.Request().Header.Get("Authorization") != token {
			return c.NoContent(http.StatusUnauthorized)
		}

		return next(c)
	}
}

func main() {
	_ = godotenv.Load()
	dataDir = os.Getenv("DATA_DIR")
	if dataDir == "" {
		log.Fatal("KEY_DIR is not defined")
	}
	token = os.Getenv("TOKEN")
	if token == "" {
		log.Fatal("TOKEN is not defined")
	}

	ConnectDbFromEnv()

	f, err := os.OpenFile(path.Join(dataDir, "startup"), syscall.O_CREAT|syscall.O_APPEND|syscall.O_WRONLY, 0644)
	if err != nil {
		log.Fatalf("KEY_DIR %s is not writable (%v)\n", dataDir, err)
	}
	_, err = f.Write([]byte(fmt.Sprintf("%s\n", time.Now().Format(TimeFormat))))
	if err != nil {
		log.Fatalf("KEY_DIR test file couldn't be written to (%v)\n", err)
	}
	err = f.Close()
	if err != nil {
		log.Fatalf("KEY_DIR test file couldn't be closed (%v)\n", err)
	}

	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)
		}
	}

	// client connection
	app.GET("/", func(c echo.Context) error {
		return c.JSON(http.StatusOK, "Hello World!")
	})
	app.GET("/keys", keys)

	// administration
	app.GET("/admin/clients", func(c echo.Context) error {
		if c.Request().Header.Get("Authorization") != token {
			return c.NoContent(http.StatusUnauthorized)
		}

		return c.JSON(http.StatusOK, clients)
	})
	app.GET("/admin/:id/logs", authed(func(c echo.Context) error {
		// TODO: finish
		return nil
	}))
	app.POST("/admin/:id/name", authed(func(c echo.Context) error {
		return nil
	}))
	app.POST("/admin/exit", authed(func(c echo.Context) error {
		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.Obj.ID == data.ID {
				client.ExitWanted = 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()

	now := time.Now()
	client := Client{
		Obj: &Dummy{
			ID:             uuid.New(),
			Name:           "Dummy",
			Connected:      true,
			ConnectedAt:    now,
			DisconnectedAt: now,
			CreatedAt:      now,
			UpdatedAt:      now,
		},
		conn:       ws,
		ExitWanted: false,
	}

	tx := db.Create(client.Obj)
	if tx.Error != nil {
		log.Error(err)
		fmt.Printf("%s client %s crashed (couldn't create record)\n", time.Now().Format(TimeFormat), client.Obj.ID.String())
		client.ExitWanted = true
		client.Obj.Connected = false
		client.Obj.DisconnectedAt = time.Now()
		return c.NoContent(http.StatusInternalServerError)
	}
	clients = append(clients, &client)
	fmt.Printf("%s client %s connected\n", time.Now().Format(TimeFormat), client.Obj.ID.String())

	err = ws.WriteMessage(websocket.TextMessage, []byte("welcome"))
	if err != nil {
		log.Error(err)
	}

	f, err := os.OpenFile(fmt.Sprintf("%s/%s_%d.txt", dataDir, time.Now().Format("2006-01-02_15-04"), client.Obj.ID.String()), syscall.O_CREAT|syscall.O_APPEND|syscall.O_WRONLY, 0644)
	if err != nil {
		log.Error(err)
		fmt.Printf("%s client %s crashed (couldn't open file)\n", time.Now().Format(TimeFormat), client.Obj.ID)
		now = time.Now()

		client.ExitWanted = true
		client.Obj.Connected = false
		client.Obj.DisconnectedAt = now
		client.Obj.UpdatedAt = now

		tx = db.Save(client.Obj)
		if tx.Error != nil {
			log.Error(err)
		}
		return c.NoContent(http.StatusInternalServerError)
	}
	defer f.Close()

	close := func() {
		_ = ws.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
		_ = ws.Close()
		now = time.Now()

		client.ExitWanted = true
		client.Obj.Connected = false
		client.Obj.DisconnectedAt = now
		client.Obj.UpdatedAt = now

		tx = db.Save(client.Obj)
		if tx.Error != nil {
			log.Error(err)
		}
	}

	for {
		if client.ExitWanted {
			close()
			fmt.Printf("%s client %s disconnected\n", time.Now().Format(TimeFormat), client.Obj.ID.String())
			break
		}

		_, msg, err := ws.ReadMessage()
		if err != nil {
			log.Error(err)
			close()
			fmt.Printf("%s client %s crashed (websocket error)\n", time.Now().Format(TimeFormat), client.Obj.ID.String())
			return c.NoContent(http.StatusInternalServerError)
		}

		_, err = f.Write(msg)
		if err != nil {
			log.Error(err)
			close()
			fmt.Printf("%s client %s crashed (file write error)\n", time.Now().Format(TimeFormat), client.Obj.ID.String())
			return c.NoContent(http.StatusInternalServerError)
		}

	}

	return c.NoContent(http.StatusOK)
}