Fixed #414

- Added sticky menu
- Optimized terminate routine for nil check
- Added test case for statistic module

src/def.go

@@ -43,7 +43,7 @@ const (
 	/* Build Constants */
 	SYSTEM_NAME       = "Zoraxy"
 	SYSTEM_VERSION    = "3.1.5"
-	DEVELOPMENT_BUILD = false /* Development: Set to false to use embedded web fs */
+	DEVELOPMENT_BUILD = true /* Development: Set to false to use embedded web fs */
 
 	/* System Constants */
 	DATABASE_PATH              = "sys.db"

src/main.go

@@ -60,19 +60,31 @@ func SetupCloseHandler() {
 func ShutdownSeq() {
 	SystemWideLogger.Println("Shutting down " + SYSTEM_NAME)
 	SystemWideLogger.Println("Closing Netstats Listener")
-	netstatBuffers.Close()
+	if netstatBuffers != nil {
+		netstatBuffers.Close()
+	}
+
 	SystemWideLogger.Println("Closing Statistic Collector")
-	statisticCollector.Close()
+	if statisticCollector != nil {
+		statisticCollector.Close()
+	}
+
 	if mdnsTickerStop != nil {
 		SystemWideLogger.Println("Stopping mDNS Discoverer (might take a few minutes)")
 		// Stop the mdns service
 		mdnsTickerStop <- true
 	}
-	mdnsScanner.Close()
+	if mdnsScanner != nil {
+		mdnsScanner.Close()
+	}
 	SystemWideLogger.Println("Shutting down load balancer")
-	loadBalancer.Close()
+	if loadBalancer != nil {
+		loadBalancer.Close()
+	}
 	SystemWideLogger.Println("Closing Certificates Auto Renewer")
-	acmeAutoRenewer.Close()
+	if acmeAutoRenewer != nil {
+		acmeAutoRenewer.Close()
+	}
 	//Remove the tmp folder
 	SystemWideLogger.Println("Cleaning up tmp files")
 	os.RemoveAll("./tmp")

src/mod/database/database.go

@@ -28,14 +28,16 @@ func NewDatabase(dbfile string, backendType dbinc.BackendType) (*Database, error
 	return newDatabase(dbfile, backendType)
 }
 
+// Get the recommended backend type for the current system
 func GetRecommendedBackendType() dbinc.BackendType {
 	//Check if the system is running on RISCV hardware
 	if runtime.GOARCH == "riscv64" {
 		//RISCV hardware, currently only support FS emulated database
 		return dbinc.BackendFSOnly
 	} else if runtime.GOOS == "windows" || (runtime.GOOS == "linux" && runtime.GOARCH == "amd64") {
-		//Powerful hardware, use LevelDB
+		//Powerful hardware
-		return dbinc.BackendLevelDB
+		return dbinc.BackendBoltDB
+		//return dbinc.BackendLevelDB
 	}
 
 	//Default to BoltDB, the safest option

src/mod/database/dbleveldb/dbleveldb.go

@@ -2,9 +2,11 @@ package dbleveldb
 
 import (
 	"encoding/json"
+	"log"
 	"path/filepath"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/syndtr/goleveldb/leveldb"
 	"github.com/syndtr/goleveldb/leveldb/util"
@@ -15,8 +17,11 @@ import (
 var _ dbinc.Backend = (*DB)(nil)
 
 type DB struct {
-	db    *leveldb.DB
+	db               *leveldb.DB
-	Table sync.Map //For emulating table creation
+	Table            sync.Map      //For emulating table creation
+	batch            leveldb.Batch //Batch write
+	writeFlushTicker *time.Ticker  //Ticker for flushing data into disk
+	writeFlushStop   chan bool     //Stop channel for write flush ticker
 }
 
 func NewDB(path string) (*DB, error) {
@@ -29,7 +34,39 @@ func NewDB(path string) (*DB, error) {
 	if err != nil {
 		return nil, err
 	}
-	return &DB{db: db, Table: sync.Map{}}, nil
+
+	thisDB := &DB{
+		db:    db,
+		Table: sync.Map{},
+		batch: leveldb.Batch{},
+	}
+
+	//Create a ticker to flush data into disk every 5 seconds
+	writeFlushTicker := time.NewTicker(5 * time.Second)
+	writeFlushStop := make(chan bool)
+	go func() {
+		for {
+			select {
+			case <-writeFlushTicker.C:
+				if thisDB.batch.Len() == 0 {
+					//No flushing needed
+					continue
+				}
+				err = db.Write(&thisDB.batch, nil)
+				if err != nil {
+					log.Println("[LevelDB] Failed to flush data into disk: ", err)
+				}
+				thisDB.batch.Reset()
+			case <-writeFlushStop:
+				return
+			}
+		}
+	}()
+
+	thisDB.writeFlushTicker = writeFlushTicker
+	thisDB.writeFlushStop = writeFlushStop
+
+	return thisDB, nil
 }
 
 func (d *DB) NewTable(tableName string) error {
@@ -66,7 +103,8 @@ func (d *DB) Write(tableName string, key string, value interface{}) error {
 	if err != nil {
 		return err
 	}
-	return d.db.Put([]byte(filepath.ToSlash(filepath.Join(tableName, key))), data, nil)
+	d.batch.Put([]byte(filepath.ToSlash(filepath.Join(tableName, key))), data)
+	return nil
 }
 
 func (d *DB) Read(tableName string, key string, assignee interface{}) error {
@@ -106,5 +144,9 @@ func (d *DB) ListTable(tableName string) ([][][]byte, error) {
 }
 
 func (d *DB) Close() {
+	//Write the remaining data in batch back into disk
+	d.writeFlushStop <- true
+	d.writeFlushTicker.Stop()
+	d.db.Write(&d.batch, nil)
 	d.db.Close()
 }

src/mod/statistic/statistic.go

@@ -33,15 +33,15 @@ type DailySummary struct {
 }
 
 type RequestInfo struct {
-	IpAddr                        string
+	IpAddr                        string //IP address of the downstream request
-	RequestOriginalCountryISOCode string
+	RequestOriginalCountryISOCode string //ISO code of the country where the request originated
-	Succ                          bool
+	Succ                          bool   //If the request is successful and resp generated by upstream instead of Zoraxy (except static web server)
-	StatusCode                    int
+	StatusCode                    int    //HTTP status code of the request
-	ForwardType                   string
+	ForwardType                   string //Forward type of the request, usually the proxy type (e.g. host-http, subdomain-websocket or vdir-http or any of the combination)
-	Referer                       string
+	Referer                       string //Referer of the downstream request
-	UserAgent                     string
+	UserAgent                     string //UserAgent of the downstream request
-	RequestURL                    string
+	RequestURL                    string //Request URL
-	Target                        string
+	Target                        string //Target domain or hostname
 }
 
 type CollectorOption struct {
@@ -59,7 +59,7 @@ func NewStatisticCollector(option CollectorOption) (*Collector, error) {
 
 	//Create the collector object
 	thisCollector := Collector{
-		DailySummary: newDailySummary(),
+		DailySummary: NewDailySummary(),
 		Option:       &option,
 	}
 
@@ -87,6 +87,11 @@ func (c *Collector) SaveSummaryOfDay() {
 	c.Option.Database.Write("stats", summaryKey, saveData)
 }
 
+// Get the daily summary up until now
+func (c *Collector) GetCurrentDailySummary() *DailySummary {
+	return c.DailySummary
+}
+
 // Load the summary of a day given
 func (c *Collector) LoadSummaryOfDay(year int, month time.Month, day int) *DailySummary {
 	date := time.Date(year, time.Month(month), day, 0, 0, 0, 0, time.Local)
@@ -99,7 +104,7 @@ func (c *Collector) LoadSummaryOfDay(year int, month time.Month, day int) *Daily
 
 // Reset today summary, for debug or restoring injections
 func (c *Collector) ResetSummaryOfDay() {
-	c.DailySummary = newDailySummary()
+	c.DailySummary = NewDailySummary()
 }
 
 // This function gives the current slot in the 288- 5 minutes interval of the day
@@ -185,8 +190,6 @@ func (c *Collector) RecordRequest(ri RequestInfo) {
 			c.DailySummary.UserAgent.Store(ri.UserAgent, ua.(int)+1)
 		}
 
-		//ADD MORE HERE IF NEEDED
-
 		//Record request URL, if it is a page
 		ext := filepath.Ext(ri.RequestURL)
 
@@ -201,6 +204,8 @@ func (c *Collector) RecordRequest(ri RequestInfo) {
 			c.DailySummary.RequestURL.Store(ri.RequestURL, ru.(int)+1)
 		}
 	}()
+
+	//ADD MORE HERE IF NEEDED
 }
 
 // nightly task
@@ -223,7 +228,7 @@ func (c *Collector) ScheduleResetRealtimeStats() chan bool {
 			case <-time.After(duration):
 				// store daily summary to database and reset summary
 				c.SaveSummaryOfDay()
-				c.DailySummary = newDailySummary()
+				c.DailySummary = NewDailySummary()
 			case <-doneCh:
 				// stop the routine
 				return
@@ -234,7 +239,7 @@ func (c *Collector) ScheduleResetRealtimeStats() chan bool {
 	return doneCh
 }
 
-func newDailySummary() *DailySummary {
+func NewDailySummary() *DailySummary {
 	return &DailySummary{
 		TotalRequest:    0,
 		ErrorRequest:    0,
@@ -247,3 +252,30 @@ func newDailySummary() *DailySummary {
 		RequestURL:      &sync.Map{},
 	}
 }
+
+func PrintDailySummary(summary *DailySummary) {
+	summary.ForwardTypes.Range(func(key, value interface{}) bool {
+		println(key.(string), value.(int))
+		return true
+	})
+	summary.RequestOrigin.Range(func(key, value interface{}) bool {
+		println(key.(string), value.(int))
+		return true
+	})
+	summary.RequestClientIp.Range(func(key, value interface{}) bool {
+		println(key.(string), value.(int))
+		return true
+	})
+	summary.Referer.Range(func(key, value interface{}) bool {
+		println(key.(string), value.(int))
+		return true
+	})
+	summary.UserAgent.Range(func(key, value interface{}) bool {
+		println(key.(string), value.(int))
+		return true
+	})
+	summary.RequestURL.Range(func(key, value interface{}) bool {
+		println(key.(string), value.(int))
+		return true
+	})
+}

src/mod/statistic/statistic_test.go

@@ -0,0 +1,215 @@
+package statistic_test
+
+import (
+	"net"
+	"os"
+	"testing"
+	"time"
+
+	"math/rand"
+
+	"imuslab.com/zoraxy/mod/database"
+	"imuslab.com/zoraxy/mod/database/dbinc"
+	"imuslab.com/zoraxy/mod/geodb"
+	"imuslab.com/zoraxy/mod/statistic"
+)
+
+const test_db_path = "test_db"
+
+func getNewDatabase() *database.Database {
+	db, err := database.NewDatabase(test_db_path, dbinc.BackendLevelDB)
+	if err != nil {
+		panic(err)
+	}
+	db.NewTable("stats")
+	return db
+}
+
+func clearDatabase(db *database.Database) {
+	db.Close()
+	os.RemoveAll(test_db_path)
+}
+
+func TestNewStatisticCollector(t *testing.T) {
+	db := getNewDatabase()
+	defer clearDatabase(db)
+	option := statistic.CollectorOption{Database: db}
+	collector, err := statistic.NewStatisticCollector(option)
+	if err != nil {
+		t.Fatalf("Expected no error, got %v", err)
+	}
+	if collector == nil {
+		t.Fatalf("Expected collector, got nil")
+	}
+
+}
+
+func TestSaveSummaryOfDay(t *testing.T) {
+	db := getNewDatabase()
+	defer clearDatabase(db)
+	option := statistic.CollectorOption{Database: db}
+	collector, _ := statistic.NewStatisticCollector(option)
+	collector.SaveSummaryOfDay()
+	// Add assertions to check if data is saved correctly
+}
+
+func TestLoadSummaryOfDay(t *testing.T) {
+	db := getNewDatabase()
+	defer clearDatabase(db)
+	option := statistic.CollectorOption{Database: db}
+	collector, _ := statistic.NewStatisticCollector(option)
+	year, month, day := time.Now().Date()
+	summary := collector.LoadSummaryOfDay(year, month, day)
+	if summary == nil {
+		t.Fatalf("Expected summary, got nil")
+	}
+}
+
+func TestResetSummaryOfDay(t *testing.T) {
+	db := getNewDatabase()
+	defer clearDatabase(db)
+	option := statistic.CollectorOption{Database: db}
+	collector, _ := statistic.NewStatisticCollector(option)
+	collector.ResetSummaryOfDay()
+	if collector.DailySummary.TotalRequest != 0 {
+		t.Fatalf("Expected TotalRequest to be 0, got %v", collector.DailySummary.TotalRequest)
+	}
+}
+
+func TestGetCurrentRealtimeStatIntervalId(t *testing.T) {
+	db := getNewDatabase()
+	defer clearDatabase(db)
+	option := statistic.CollectorOption{Database: db}
+	collector, _ := statistic.NewStatisticCollector(option)
+	intervalId := collector.GetCurrentRealtimeStatIntervalId()
+	if intervalId < 0 || intervalId > 287 {
+		t.Fatalf("Expected intervalId to be between 0 and 287, got %v", intervalId)
+	}
+}
+
+func TestRecordRequest(t *testing.T) {
+	db := getNewDatabase()
+	defer clearDatabase(db)
+	option := statistic.CollectorOption{Database: db}
+	collector, _ := statistic.NewStatisticCollector(option)
+	requestInfo := statistic.RequestInfo{
+		IpAddr:                        "127.0.0.1",
+		RequestOriginalCountryISOCode: "US",
+		Succ:                          true,
+		StatusCode:                    200,
+		ForwardType:                   "type1",
+		Referer:                       "http://example.com",
+		UserAgent:                     "Mozilla/5.0",
+		RequestURL:                    "/test",
+		Target:                        "target1",
+	}
+	collector.RecordRequest(requestInfo)
+	time.Sleep(1 * time.Second) // Wait for the goroutine to finish
+	if collector.DailySummary.TotalRequest != 1 {
+		t.Fatalf("Expected TotalRequest to be 1, got %v", collector.DailySummary.TotalRequest)
+	}
+}
+
+func TestScheduleResetRealtimeStats(t *testing.T) {
+	db := getNewDatabase()
+	defer clearDatabase(db)
+	option := statistic.CollectorOption{Database: db}
+	collector, _ := statistic.NewStatisticCollector(option)
+	stopChan := collector.ScheduleResetRealtimeStats()
+	if stopChan == nil {
+		t.Fatalf("Expected stopChan, got nil")
+	}
+	collector.Close()
+}
+
+func TestNewDailySummary(t *testing.T) {
+	summary := statistic.NewDailySummary()
+	if summary.TotalRequest != 0 {
+		t.Fatalf("Expected TotalRequest to be 0, got %v", summary.TotalRequest)
+	}
+	if summary.ForwardTypes == nil {
+		t.Fatalf("Expected ForwardTypes to be initialized, got nil")
+	}
+	if summary.RequestOrigin == nil {
+		t.Fatalf("Expected RequestOrigin to be initialized, got nil")
+	}
+	if summary.RequestClientIp == nil {
+		t.Fatalf("Expected RequestClientIp to be initialized, got nil")
+	}
+	if summary.Referer == nil {
+		t.Fatalf("Expected Referer to be initialized, got nil")
+	}
+	if summary.UserAgent == nil {
+		t.Fatalf("Expected UserAgent to be initialized, got nil")
+	}
+	if summary.RequestURL == nil {
+		t.Fatalf("Expected RequestURL to be initialized, got nil")
+	}
+}
+
+func generateTestRequestInfo(db *database.Database) statistic.RequestInfo {
+	//Generate a random IPv4 address
+	randomIpAddr := ""
+	for {
+		ip := net.IPv4(byte(rand.Intn(256)), byte(rand.Intn(256)), byte(rand.Intn(256)), byte(rand.Intn(256)))
+		if !ip.IsPrivate() && !ip.IsLoopback() && !ip.IsMulticast() && !ip.IsUnspecified() {
+			randomIpAddr = ip.String()
+			break
+		}
+	}
+
+	//Resolve the country code for this IP
+	ipLocation := "unknown"
+	geoIpResolver, err := geodb.NewGeoDb(db, &geodb.StoreOptions{
+		AllowSlowIpv4LookUp: false,
+		AllowSlowIpv6Lookup: true, //Just to save some RAM
+	})
+
+	if err == nil {
+		ipInfo, _ := geoIpResolver.ResolveCountryCodeFromIP(randomIpAddr)
+		ipLocation = ipInfo.CountryIsoCode
+	}
+
+	forwardType := "host-http"
+	//Generate a random forward type between "subdomain-http" and "host-https"
+	if rand.Intn(2) == 1 {
+		forwardType = "subdomain-http"
+	}
+
+	//Generate 5 random refers URL and pick from there
+	referers := []string{"https://example.com", "https://example.org", "https://example.net", "https://example.io", "https://example.co"}
+	referer := referers[rand.Intn(5)]
+
+	return statistic.RequestInfo{
+		IpAddr:                        randomIpAddr,
+		RequestOriginalCountryISOCode: ipLocation,
+		Succ:                          true,
+		StatusCode:                    200,
+		ForwardType:                   forwardType,
+		Referer:                       referer,
+		UserAgent:                     "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36",
+		RequestURL:                    "/benchmark",
+		Target:                        "test.imuslab.internal",
+	}
+}
+
+func BenchmarkRecordRequest(b *testing.B) {
+	db := getNewDatabase()
+	defer clearDatabase(db)
+
+	option := statistic.CollectorOption{Database: db}
+	collector, _ := statistic.NewStatisticCollector(option)
+	var requestInfo statistic.RequestInfo = generateTestRequestInfo(db)
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		collector.RecordRequest(requestInfo)
+		collector.SaveSummaryOfDay()
+	}
+
+	//Write the current in-memory summary to database file
+	b.StopTimer()
+
+	//Print the generated summary
+	//testSummary := collector.GetCurrentDailySummary()
+	//statistic.PrintDailySummary(testSummary)
+}

src/web/main.css

@@ -65,6 +65,8 @@ body{
     height: calc(100% - 51px);
     overflow-y: auto;
     width: 240px;
+    position: sticky;
+    top: 4em;
 }
 
 .contentWindow{