diff --git a/src/def.go b/src/def.go index 10a940e..3b85cb7 100644 --- a/src/def.go +++ b/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" diff --git a/src/main.go b/src/main.go index e114930..580fff9 100644 --- a/src/main.go +++ b/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") diff --git a/src/mod/database/database.go b/src/mod/database/database.go index d348d15..cb831fe 100644 --- a/src/mod/database/database.go +++ b/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 - return dbinc.BackendLevelDB + //Powerful hardware + return dbinc.BackendBoltDB + //return dbinc.BackendLevelDB } //Default to BoltDB, the safest option diff --git a/src/mod/database/dbleveldb/dbleveldb.go b/src/mod/database/dbleveldb/dbleveldb.go index 91b9155..2511bcf 100644 --- a/src/mod/database/dbleveldb/dbleveldb.go +++ b/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 - Table sync.Map //For emulating table creation + db *leveldb.DB + 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() } diff --git a/src/mod/statistic/statistic.go b/src/mod/statistic/statistic.go index 06bf11b..eb58175 100644 --- a/src/mod/statistic/statistic.go +++ b/src/mod/statistic/statistic.go @@ -33,15 +33,15 @@ type DailySummary struct { } type RequestInfo struct { - IpAddr string - RequestOriginalCountryISOCode string - Succ bool - StatusCode int - ForwardType string - Referer string - UserAgent string - RequestURL string - Target string + IpAddr string //IP address of the downstream request + RequestOriginalCountryISOCode string //ISO code of the country where the request originated + Succ bool //If the request is successful and resp generated by upstream instead of Zoraxy (except static web server) + StatusCode int //HTTP status code of the request + 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 of the downstream request + UserAgent string //UserAgent of the downstream request + RequestURL string //Request URL + 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 + }) +} diff --git a/src/mod/statistic/statistic_test.go b/src/mod/statistic/statistic_test.go new file mode 100644 index 0000000..97da9da --- /dev/null +++ b/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) +} diff --git a/src/web/main.css b/src/web/main.css index e973898..3989f5b 100644 --- a/src/web/main.css +++ b/src/web/main.css @@ -65,6 +65,8 @@ body{ height: calc(100% - 51px); overflow-y: auto; width: 240px; + position: sticky; + top: 4em; } .contentWindow{