server.go 7.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278
  1. package server
  2. import (
  3. "context"
  4. "fmt"
  5. "io"
  6. "log"
  7. "net"
  8. "net/http"
  9. "os"
  10. "os/signal"
  11. "path/filepath"
  12. "strings"
  13. "sync"
  14. "github.com/claudiodangelis/qrcp/config"
  15. "github.com/claudiodangelis/qrcp/pages"
  16. "github.com/claudiodangelis/qrcp/payload"
  17. "github.com/claudiodangelis/qrcp/util"
  18. "gopkg.in/cheggaaa/pb.v1"
  19. )
  20. // Server is the server
  21. type Server struct {
  22. // SendURL is the URL used to send the file
  23. SendURL string
  24. // ReceiveURL is the URL used to Receive the file
  25. ReceiveURL string
  26. instance *http.Server
  27. payload payload.Payload
  28. outputDir string
  29. stopChannel chan bool
  30. // expectParallelRequests is set to true when qrcp sends files, in order
  31. // to support downloading of parallel chunks
  32. expectParallelRequests bool
  33. }
  34. // ReceiveTo sets the output directory
  35. func (s *Server) ReceiveTo(dir string) error {
  36. output, err := filepath.Abs(dir)
  37. if err != nil {
  38. return err
  39. }
  40. // Check if the output dir exists
  41. fileinfo, err := os.Stat(output)
  42. if err != nil {
  43. return err
  44. }
  45. if !fileinfo.IsDir() {
  46. return fmt.Errorf("%s is not a valid directory", output)
  47. }
  48. s.outputDir = output
  49. return nil
  50. }
  51. // Send adds a handler for sending the file
  52. func (s *Server) Send(p payload.Payload) {
  53. s.payload = p
  54. s.expectParallelRequests = true
  55. }
  56. // Wait for transfer to be completed, it waits forever if kept awlive
  57. func (s Server) Wait() error {
  58. <-s.stopChannel
  59. if err := s.instance.Shutdown(context.Background()); err != nil {
  60. log.Println(err)
  61. }
  62. if s.payload.DeleteAfterTransfer {
  63. s.payload.Delete()
  64. }
  65. return nil
  66. }
  67. // New instance of the server
  68. func New(cfg *config.Config) (*Server, error) {
  69. app := &Server{}
  70. // Get the address of the configured interface to bind the server to
  71. bind, err := util.GetInterfaceAddress(cfg.Interface)
  72. if err != nil {
  73. return &Server{}, err
  74. }
  75. // Create a listener. If `port: 0`, a random one is chosen
  76. listener, err := net.Listen("tcp", fmt.Sprintf("%s:%d", bind, cfg.Port))
  77. if err != nil {
  78. return nil, err
  79. }
  80. // Set the value of computed port
  81. port := listener.Addr().(*net.TCPAddr).Port
  82. // Set the host
  83. host := fmt.Sprintf("%s:%d", bind, port)
  84. // Get a random path to use
  85. path := cfg.Path
  86. if path == "" {
  87. path = util.GetRandomURLPath()
  88. }
  89. // Set the hostname
  90. hostname := fmt.Sprintf("%s:%d", bind, port)
  91. // Use external IP when using `interface: any`, unless a FQDN is set
  92. if bind == "0.0.0.0" && cfg.FQDN == "" {
  93. fmt.Println("Retrieving the external IP...")
  94. extIP, err := util.GetExernalIP()
  95. if err != nil {
  96. panic(err)
  97. }
  98. hostname = fmt.Sprintf("%s:%d", extIP.String(), port)
  99. }
  100. // Use a fully-qualified domain name if set
  101. if cfg.FQDN != "" {
  102. hostname = fmt.Sprintf("%s:%d", cfg.FQDN, port)
  103. }
  104. // Set send and receive URLs
  105. app.SendURL = fmt.Sprintf("http://%s/send/%s",
  106. hostname, path)
  107. app.ReceiveURL = fmt.Sprintf("http://%s/receive/%s",
  108. hostname, path)
  109. // Create a server
  110. httpserver := &http.Server{Addr: host}
  111. // Create channel to send message to stop server
  112. app.stopChannel = make(chan bool)
  113. // Create cookie used to verify request is coming from first client to connect
  114. cookie := http.Cookie{Name: "qrcp", Value: ""}
  115. // Gracefully shutdown when an OS signal is received
  116. sig := make(chan os.Signal, 1)
  117. signal.Notify(sig, os.Interrupt)
  118. go func() {
  119. <-sig
  120. app.stopChannel <- true
  121. }()
  122. // The handler adds and removes from the sync.WaitGroup
  123. // When the group is zero all requests are completed
  124. // and the server is shutdown
  125. var waitgroup sync.WaitGroup
  126. waitgroup.Add(1)
  127. var initCookie sync.Once
  128. // Create handlers
  129. // Send handler (sends file to caller)
  130. http.HandleFunc("/send/"+path, func(w http.ResponseWriter, r *http.Request) {
  131. if cookie.Value == "" {
  132. if !strings.HasPrefix(r.Header.Get("User-Agent"), "Mozilla") {
  133. http.Error(w, "", http.StatusOK)
  134. return
  135. }
  136. initCookie.Do(func() {
  137. value, err := util.GetSessionID()
  138. if err != nil {
  139. log.Println("Unable to generate session ID", err)
  140. app.stopChannel <- true
  141. return
  142. }
  143. cookie.Value = value
  144. http.SetCookie(w, &cookie)
  145. })
  146. } else {
  147. // Check for the expected cookie and value
  148. // If it is missing or doesn't match
  149. // return a 404 status
  150. rcookie, err := r.Cookie(cookie.Name)
  151. if err != nil || rcookie.Value != cookie.Value {
  152. http.Error(w, "", http.StatusNotFound)
  153. return
  154. }
  155. // If the cookie exits and matches
  156. // this is an aadditional request.
  157. // Increment the waitgroup
  158. waitgroup.Add(1)
  159. }
  160. // Remove connection from the waitfroup when done
  161. defer waitgroup.Done()
  162. w.Header().Set("Content-Disposition", "attachment; filename="+
  163. app.payload.Filename)
  164. http.ServeFile(w, r, app.payload.Path)
  165. })
  166. // Upload handler (serves the upload page)
  167. http.HandleFunc("/receive/"+path, func(w http.ResponseWriter, r *http.Request) {
  168. htmlVariables := struct {
  169. Route string
  170. File string
  171. }{}
  172. htmlVariables.Route = "/receive/" + path
  173. switch r.Method {
  174. case "POST":
  175. filenames := util.ReadFilenames(app.outputDir)
  176. reader, err := r.MultipartReader()
  177. if err != nil {
  178. fmt.Fprintf(w, "Upload error: %v\n", err)
  179. log.Printf("Upload error: %v\n", err)
  180. app.stopChannel <- true
  181. return
  182. }
  183. transferredFiles := []string{}
  184. progressBar := pb.New64(r.ContentLength)
  185. progressBar.ShowCounters = false
  186. for {
  187. part, err := reader.NextPart()
  188. if err == io.EOF {
  189. break
  190. }
  191. // iIf part.FileName() is empty, skip this iteration.
  192. if part.FileName() == "" {
  193. continue
  194. }
  195. // Prepare the destination
  196. fileName := getFileName(part.FileName(), filenames)
  197. out, err := os.Create(filepath.Join(app.outputDir, fileName))
  198. if err != nil {
  199. // Output to server
  200. fmt.Fprintf(w, "Unable to create the file for writing: %s\n", err)
  201. // Output to console
  202. log.Printf("Unable to create the file for writing: %s\n", err)
  203. // Send signal to server to shutdown
  204. app.stopChannel <- true
  205. return
  206. }
  207. defer out.Close()
  208. // Add name of new file
  209. filenames = append(filenames, fileName)
  210. // Write the content from POSTed file to the out
  211. fmt.Println("Transferring file: ", out.Name())
  212. progressBar.Prefix(out.Name())
  213. progressBar.Start()
  214. buf := make([]byte, 1024)
  215. for {
  216. // Read a chunk
  217. n, err := part.Read(buf)
  218. if err != nil && err != io.EOF {
  219. // Output to server
  220. fmt.Fprintf(w, "Unable to write file to disk: %v", err)
  221. // Output to console
  222. fmt.Printf("Unable to write file to disk: %v", err)
  223. // Send signal to server to shutdown
  224. app.stopChannel <- true
  225. return
  226. }
  227. if n == 0 {
  228. break
  229. }
  230. // Write a chunk
  231. if _, err := out.Write(buf[:n]); err != nil {
  232. // Output to server
  233. fmt.Fprintf(w, "Unable to write file to disk: %v", err)
  234. // Output to console
  235. log.Printf("Unable to write file to disk: %v", err)
  236. // Send signal to server to shutdown
  237. app.stopChannel <- true
  238. return
  239. }
  240. progressBar.Add(n)
  241. }
  242. transferredFiles = append(transferredFiles, out.Name())
  243. }
  244. progressBar.FinishPrint("File transfer completed")
  245. // Set the value of the variable to the actually transferred files
  246. htmlVariables.File = strings.Join(transferredFiles, ", ")
  247. serveTemplate("done", pages.Done, w, htmlVariables)
  248. if cfg.KeepAlive == false {
  249. app.stopChannel <- true
  250. }
  251. case "GET":
  252. serveTemplate("upload", pages.Upload, w, htmlVariables)
  253. }
  254. })
  255. // Wait for all wg to be done, then send shutdown signal
  256. go func() {
  257. waitgroup.Wait()
  258. if cfg.KeepAlive || !app.expectParallelRequests {
  259. return
  260. }
  261. app.stopChannel <- true
  262. }()
  263. // Receive handler (receives file from caller)
  264. go func() {
  265. if err := (httpserver.Serve(tcpKeepAliveListener{listener.(*net.TCPListener)})); err != http.ErrServerClosed {
  266. log.Fatalln(err)
  267. }
  268. }()
  269. app.instance = httpserver
  270. return app, nil
  271. }