From 2c3532f12f8900f736319de7446f2fef0a8117f4 Mon Sep 17 00:00:00 2001
From: heimoshuiyu <heimoshuiyu@gmail.com>
Date: Thu, 30 Nov 2023 10:18:01 +0800
Subject: [PATCH] add hostname

---
 main.go      | 12 ++++++++++++
 record.go    |  1 +
 structure.go |  1 +
 3 files changed, 14 insertions(+)

diff --git a/main.go b/main.go
index 36c56b1..4a470ae 100644
--- a/main.go
+++ b/main.go
@@ -4,6 +4,7 @@ import (
 	"flag"
 	"fmt"
 	"log"
+	"os"
 	"strings"
 	"time"
 
@@ -38,11 +39,17 @@ func main() {
 			PrepareStmt:            true,
 			SkipDefaultTransaction: true,
 		})
+		if err != nil {
+			log.Fatalf("Error to connect sqlite database: %s", err)
+		}
 	case "postgres":
 		db, err = gorm.Open(postgres.Open(config.DBAddr), &gorm.Config{
 			PrepareStmt:            true,
 			SkipDefaultTransaction: true,
 		})
+		if err != nil {
+			log.Fatalf("Error to connect postgres database: %s", err)
+		}
 	default:
 		log.Fatalf("Unsupported database type: '%s'", config.DBType)
 	}
@@ -88,8 +95,13 @@ func main() {
 	})
 
 	engine.POST("/v1/*any", func(c *gin.Context) {
+		hostname, err := os.Hostname()
+		if config.Hostname != "" {
+			hostname = config.Hostname
+		}
 		record := Record{
 			IP:            c.ClientIP(),
+			Hostname:      hostname,
 			CreatedAt:     time.Now(),
 			Authorization: c.Request.Header.Get("Authorization"),
 			UserAgent:     c.Request.Header.Get("User-Agent"),
diff --git a/record.go b/record.go
index 5f8d6a8..58b053f 100644
--- a/record.go
+++ b/record.go
@@ -6,6 +6,7 @@ import (
 
 type Record struct {
 	ID               int64 `gorm:"primaryKey,autoIncrement"`
+	Hostname         string
 	UpstreamEndpoint string
 	UpstreamSK       string
 	CreatedAt        time.Time
diff --git a/structure.go b/structure.go
index ad8037d..6f48bae 100644
--- a/structure.go
+++ b/structure.go
@@ -9,6 +9,7 @@ import (
 
 type Config struct {
 	Address       string            `yaml:"address"`
+	Hostname      string            `yaml:"hostname"`
 	DBType        string            `yaml:"dbtype"`
 	DBAddr        string            `yaml:"dbaddr"`
 	Authorization string            `yaml:"authorization"`