server/main.go

package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"os"
	"strings"

	"gopkg.in/ini.v1"
	"gorm.io/driver/postgres"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
)

// Global variables
// DB global database oject
var DB *gorm.DB

func handleInitialConnection(conn net.Conn) {
	initial_request := make([]byte, 256)

	n, err := conn.Read(initial_request)
	if err == io.EOF {
		fmt.Fprintf(os.Stderr, "ERROR: client has disconnected: %s\n", conn.RemoteAddr())
		conn.Close()
		return
	} else if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %s failed with the error %v, closing connection\n", conn.RemoteAddr(), err)
		conn.Close()
		return
	}

	if n == 0 {
		fmt.Fprintf(os.Stderr, "ERROR: received zero bytes, closing connection\n")
		conn.Close()
		return
	}

	data := Data{}
	err = json.Unmarshal(initial_request[:n], &data)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: failed to parse data in initial request, closing connection\n")
		conn.Close()
		return
	}

	if *data.Type != REQUEST && (*data.Request != NEW_SERVER_CLIENT || *data.Request != NEW_DESKTOP_CLIENT) { // TODO: I think this logic may be a little wrong
		fmt.Fprintf(os.Stderr, "ERROR: invalid initial request, closing connection from %s\n", conn.RemoteAddr())
		conn.Close()
		return
	}

	switch *data.Request {
	case NEW_SERVER_CLIENT:
		var connectedResponse Data
		connectedResponse.Type = new(Type)
		connectedResponse.Response = new(Response_Type)
		*connectedResponse.Type = RESPONSE
		foundHost := false
		for _, host := range GetAllHosts() {
			if strings.ToLower(host.Hostname) == strings.ToLower(*data.Hostname) {
				*connectedResponse.Type = RESPONSE
				*connectedResponse.Response = OK

				fmt.Printf("Found host in database\n")

				foundHost = true

				break
			}
		}

		if !foundHost {
			if ok := CreateHost(*data.Hostname); !ok {
				fmt.Fprintf(os.Stderr, "ERROR: failed to create host in database, disconnecting client\n")
				conn.Close()
				return
			}
			fmt.Printf("Created host in database\n")
		}

		bytes, err := json.Marshal(connectedResponse)
		if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: failed to create JSON to send back to client, closing connection\n")
			conn.Close()
			return
		}

		fmt.Println("Response: ", string(bytes))

		conn.Write(bytes)

		// TODO: Handle client connections differently
		handleServerConnection(conn)
	}
}

func handleServerConnection(conn net.Conn) {
	data := make([]byte, 4028)
	for {
		n, err := conn.Read(data)
		if err == io.EOF {
			fmt.Fprintf(os.Stderr, "ERROR: client has disconnected: %s\n", conn.RemoteAddr())
			conn.Close()
			return
		} else if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: %s failed with the error %v, closing connection\n", conn.RemoteAddr(), err)
			conn.Close()
			return
		}

		if n == 0 {
			continue
		}

		fmt.Printf("Read %d bytes from %s\n", n, conn.RemoteAddr())
		fmt.Printf("data:\n\t'%s'\n", data)
	}
}

func main() {
	cfg, err := ini.Load("config.ini")
	if err != nil {
		log.Fatal("Fail to read file: ", err)
	}

	// Classic read of values, default section can be represented as empty string
	log.Println("App Mode:         ", cfg.Section("").Key("app_mode").String())
	log.Println("Database Type:    ", cfg.Section("database").Key("type").String())

	if strings.ToLower(cfg.Section("database").Key("type").String()) == "postgres" {
		log.Println("Database Host:    ", cfg.Section("database").Key("host").String())
		log.Println("Database Port:    ", cfg.Section("database").Key("port").String())
		log.Println("Database Username:", cfg.Section("database").Key("username").String())
		log.Println("Database Name:    ", cfg.Section("database").Key("DB_name").String())

		databaseConnectString := fmt.Sprintf(
			"host=%s port=%s user=%s DBname=%s password=%s sslmode=disable",
			cfg.Section("database").Key("host"),
			cfg.Section("database").Key("port"),
			cfg.Section("database").Key("user"),
			cfg.Section("database").Key("DB_name"),
			cfg.Section("database").Key("password"),
		)

		log.Println(databaseConnectString)
		DB, err = gorm.Open(postgres.Open(databaseConnectString), &gorm.Config{TranslateError: true})
		if err != nil {
			log.Fatal("Failed to connect to PostgreSQL DB")
		}
	} else if strings.ToLower(cfg.Section("database").Key("type").String()) == "sqlite" {
		sqlite_path := cfg.Section("database").Key("path").String()
		log.Println("sqlite DB Path:    ", sqlite_path)

		DB, err = gorm.Open(sqlite.Open(sqlite_path), &gorm.Config{TranslateError: true})
		if err != nil {
			log.Fatal("ERROR: failed to open sqlite database")
		}
	}

	port := cfg.Section("connection").Key("port").String()

	log.Println("Successfully connected to database")

	err = DB.AutoMigrate(&CPU{})
	if err != nil {
		log.Fatal("Failed to migrate CPU schema")
	}
	err = DB.AutoMigrate(&Memory{})
	if err != nil {
		log.Fatal("Failed to migrate Memory schema")
	}
	err = DB.AutoMigrate(&Host{})
	if err != nil {
		log.Fatal("Failed to migrate Host schema")
	}

	addr := fmt.Sprintf(":%s", port)
	ln, err := net.Listen("tcp", addr)
	if err != nil {
		fmt.Fprintln(os.Stderr, "ERROR: failed to listen on port with error: ", err)
		os.Exit(-1)
	}
	fmt.Printf("Listening on %s\n", addr)
	for {
		conn, err := ln.Accept()
		if err != nil {
			fmt.Fprintln(os.Stderr, "ERROR: failed to connect the error: ", err)
			continue
		}

		go handleInitialConnection(conn)
	}
}