start.go 4.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167
  1. package main
  2. import (
  3. "bytes"
  4. "fmt"
  5. "io"
  6. "io/fs"
  7. "log"
  8. "net/http"
  9. "os"
  10. "github.com/google/uuid"
  11. "github.com/gorilla/csrf"
  12. "imuslab.com/bokofs/bokofsd/mod/disktool/raid"
  13. "imuslab.com/bokofs/bokofsd/mod/netstat"
  14. )
  15. /*
  16. start.go
  17. This file handles the startup and initialization of the application
  18. */
  19. func initialization() error {
  20. /* Check and generate system UUID */
  21. configFolderPath := "./config"
  22. if _, err := os.Stat(configFolderPath); os.IsNotExist(err) {
  23. fmt.Printf("Config folder does not exist. Creating folder at %s\n", configFolderPath)
  24. if err := os.Mkdir(configFolderPath, os.ModePerm); err != nil {
  25. return fmt.Errorf("error creating config folder: %v", err)
  26. }
  27. }
  28. // Check if sys.uuid exists, if not generate a unique UUID and write it to sys.uuid
  29. uuidFilePath := configFolderPath + "/sys.uuid"
  30. if _, err := os.Stat(uuidFilePath); os.IsNotExist(err) {
  31. newUUID := uuid.New().String()
  32. if err := os.WriteFile(uuidFilePath, []byte(newUUID), 0644); err != nil {
  33. return fmt.Errorf("error writing UUID to file: %v", err)
  34. }
  35. }
  36. // Read the UUID from sys.uuid
  37. uuidBytes, err := os.ReadFile(uuidFilePath)
  38. if err != nil {
  39. return fmt.Errorf("error reading UUID from file: %v", err)
  40. }
  41. sysuuid = string(uuidBytes)
  42. /* File system handler */
  43. if *devMode {
  44. fmt.Println("Development mode enabled. Serving files from ./web directory.")
  45. webfs = http.Dir("./web")
  46. } else {
  47. fmt.Println("Production mode enabled. Serving files from embedded filesystem.")
  48. subFS, err := fs.Sub(embeddedFiles, "web")
  49. if err != nil {
  50. fmt.Fprintf(os.Stderr, "Error accessing embedded subdirectory: %v\n", err)
  51. os.Exit(1)
  52. }
  53. webfs = http.FS(subFS)
  54. }
  55. /* Network statistics */
  56. nsb, err := netstat.NewNetStatBuffer(300)
  57. if err != nil {
  58. return fmt.Errorf("error creating netstat buffer: %v", err)
  59. }
  60. netstatBuffer = nsb
  61. /* Package Check */
  62. if !checkRuntimeEnvironment() {
  63. return fmt.Errorf("runtime environment check failed")
  64. }
  65. /* RAID Manager */
  66. rm, err := raid.NewRaidManager()
  67. if err != nil {
  68. return err
  69. }
  70. raidManager = rm
  71. /* CSRF Middleware */
  72. csrfMiddleware = csrf.Protect(
  73. []byte(sysuuid),
  74. csrf.CookieName(CSRF_COOKIENAME),
  75. csrf.Secure(false),
  76. csrf.Path("/"),
  77. csrf.SameSite(csrf.SameSiteLaxMode),
  78. )
  79. return nil
  80. }
  81. // tmplateMiddleware is a middleware that serves HTML files and injects the CSRF token
  82. func tmplMiddleware(next http.Handler) http.Handler {
  83. return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
  84. csrfToken := csrf.Token(r)
  85. // Check if the request is for a path or ends with .html
  86. if r.URL.Path == "/" || r.URL.Path[len(r.URL.Path)-5:] == ".html" {
  87. file, err := webfs.Open(r.URL.Path)
  88. if err != nil {
  89. http.NotFound(w, r)
  90. return
  91. }
  92. defer file.Close()
  93. // Check if the file is a directory
  94. fileInfo, err := file.Stat()
  95. if err != nil {
  96. log.Println(err)
  97. http.Error(w, "Error retrieving file information", http.StatusInternalServerError)
  98. return
  99. }
  100. if fileInfo.IsDir() {
  101. // If the file is a directory, try to open /index.html
  102. indexFile, err := webfs.Open(r.URL.Path + "/index.html")
  103. if err != nil {
  104. http.NotFound(w, r)
  105. return
  106. }
  107. defer indexFile.Close()
  108. file = indexFile
  109. }
  110. // Replace {{.csrfToken}} in the HTML file with the CSRF token
  111. content, err := io.ReadAll(file)
  112. if err != nil {
  113. log.Println(err)
  114. http.Error(w, "Error reading file content", http.StatusInternalServerError)
  115. return
  116. }
  117. // Replace {{.csrfToken}} with the actual CSRF token
  118. modifiedContent := bytes.Replace(content, []byte("{{.csrfToken}}"), []byte(csrfToken), -1)
  119. // Write the modified content to the response
  120. w.Header().Set("Content-Type", "text/html; charset=utf-8")
  121. w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
  122. w.Header().Set("Pragma", "no-cache")
  123. w.Header().Set("Expires", "0")
  124. w.WriteHeader(http.StatusOK)
  125. w.Write(modifiedContent)
  126. return
  127. }
  128. next.ServeHTTP(w, r)
  129. // Add template engine initialization here if needed
  130. })
  131. }
  132. // Cleanup function to be called on exit
  133. func cleanup() {
  134. fmt.Println("Performing cleanup tasks...")
  135. // Close the netstat buffer if it was initialized
  136. if netstatBuffer != nil {
  137. fmt.Println("Closing netstat buffer...")
  138. netstatBuffer.Close()
  139. }
  140. fmt.Println("Cleanup completed.")
  141. }