- Added sticky menu
- Optimized terminate routine for nil check
- Added test case for statistic module
This commit is contained in:
Toby Chui 2024-12-08 12:54:50 +08:00
parent cc08c704de
commit 8ff51044bb
7 changed files with 332 additions and 27 deletions

View File

@ -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"

View File

@ -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")

View File

@ -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

View File

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

View File

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

View File

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

View File

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