PageRenderTime 54ms CodeModel.GetById 23ms RepoModel.GetById 0ms app.codeStats 0ms

/bitsrv/main.go

https://gitlab.com/magan/gobits
Go | 441 lines | 308 code | 76 blank | 57 comment | 88 complexity | ea6afdc82e7017e08488e50a482dfd54 MD5 | raw file
  1. /*
  2. GoBITS - A server implementation of Microsoft BITS (Background Intelligent Transfer Service) written in go.
  3. Copyright (C) 2015 Magnus Andersson
  4. */
  5. package bitsrv
  6. import (
  7. "crypto/rand"
  8. "errors"
  9. "fmt"
  10. "io"
  11. "io/ioutil"
  12. "net/http"
  13. "os"
  14. "path"
  15. "path/filepath"
  16. "regexp"
  17. "strconv"
  18. "strings"
  19. )
  20. type Config struct {
  21. // Directory to store unfinished files in
  22. TempDir string
  23. // Allowed method name
  24. AllowedMethod string
  25. // Protocol to use
  26. Protocol string
  27. // Max size of uploaded file
  28. MaxSize uint64
  29. // Whitelisted filter
  30. Allowed []string
  31. // Blacklisted filter
  32. Disallowed []string
  33. }
  34. type BITSHandler struct {
  35. cfg Config
  36. callback CallbackFunc
  37. }
  38. type BITSEvent int
  39. const (
  40. EventCreateSession BITSEvent = iota
  41. EventRecieveFile
  42. EventCloseSession
  43. EventCancelSession
  44. )
  45. type CallbackFunc func(event BITSEvent, Session, Path string)
  46. const (
  47. BG_ERROR_CONTEXT_NONE = 0
  48. BG_ERROR_CONTEXT_UNKNOWN = 1
  49. BG_ERROR_CONTEXT_GENERAL_QUEUE_MANAGER = 2
  50. BG_ERROR_CONTEXT_QUEUE_MANAGER_NOTIFICATION = 3
  51. BG_ERROR_CONTEXT_LOCAL_FILE = 4
  52. BG_ERROR_CONTEXT_REMOTE_FILE = 5
  53. BG_ERROR_CONTEXT_GENERAL_TRANSPORT = 6
  54. BG_ERROR_CONTEXT_REMOTE_APPLICATION = 7
  55. )
  56. func NewHandler(cfg Config, cb CallbackFunc) (b *BITSHandler) {
  57. b = new(BITSHandler)
  58. b.cfg = cfg
  59. b.callback = cb
  60. // Set defaults
  61. if b.cfg.AllowedMethod == "" {
  62. b.cfg.AllowedMethod = "BITS_POST"
  63. }
  64. if b.cfg.Protocol == "" {
  65. // BITS 1.5 Upload Protocol
  66. // https://msdn.microsoft.com/en-us/library/aa362833(v=vs.85).aspx
  67. b.cfg.Protocol = "{7df0354d-249b-430f-820d-3d2a9bef4931}"
  68. }
  69. if b.cfg.TempDir == "" {
  70. b.cfg.TempDir = path.Join(os.TempDir(), "gobits")
  71. }
  72. if len(b.cfg.Allowed) == 0 {
  73. b.cfg.Allowed = []string{".*"}
  74. }
  75. return
  76. }
  77. func (b *BITSHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
  78. // Only allow BITS requests
  79. if r.Method != b.cfg.AllowedMethod {
  80. http.Error(w, "Unauthorized", http.StatusBadRequest)
  81. return
  82. }
  83. packetType := strings.ToLower(r.Header.Get("BITS-Packet-Type"))
  84. sessionId := r.Header.Get("BITS-Session-Id")
  85. switch packetType {
  86. case "ping":
  87. b.bitsPing(w, r)
  88. case "create-session":
  89. b.bitsCreate(w, r)
  90. case "cancel-session":
  91. b.bitsCancel(w, r, sessionId)
  92. case "close-session":
  93. b.bitsClose(w, r, sessionId)
  94. case "fragment":
  95. b.bitsFragment(w, r, sessionId)
  96. default:
  97. bitsError(w, "", http.StatusBadRequest, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  98. }
  99. }
  100. func bitsError(w http.ResponseWriter, uuid string, status, code, context int) {
  101. w.Header().Add("BITS-Packet-Type", "Ack")
  102. if uuid != "" {
  103. w.Header().Add("BITS-Session-Id", uuid)
  104. }
  105. w.Header().Add("BITS-Error-Code", strconv.FormatInt(int64(code), 16))
  106. w.Header().Add("BITS-Error-Context", strconv.FormatInt(int64(context), 16))
  107. w.WriteHeader(status)
  108. w.Write(nil)
  109. }
  110. func (b *BITSHandler) bitsPing(w http.ResponseWriter, r *http.Request) {
  111. w.Header().Add("BITS-Packet-Type", "Ack")
  112. w.Write(nil)
  113. }
  114. func (b *BITSHandler) bitsCreate(w http.ResponseWriter, r *http.Request) {
  115. // Check for correct protocol
  116. var protocol string
  117. protocols := strings.Split(r.Header.Get("BITS-Supported-Protocols"), " ")
  118. for _, protocol = range protocols {
  119. if protocol == b.cfg.AllowedMethod {
  120. break
  121. }
  122. }
  123. if protocol != b.cfg.Protocol {
  124. bitsError(w, "", http.StatusBadRequest, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  125. return
  126. }
  127. // Create new session UUID
  128. uuid, err := newUUID()
  129. if err != nil {
  130. bitsError(w, "", http.StatusInternalServerError, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  131. return
  132. }
  133. // Create session directory
  134. tmpDir := path.Join(b.cfg.TempDir, uuid)
  135. if err = os.MkdirAll(tmpDir, 0600); err != nil {
  136. // Handle error
  137. }
  138. b.callback(EventCreateSession, uuid, tmpDir)
  139. w.Header().Add("BITS-Packet-Type", "Ack")
  140. w.Header().Add("BITS-Protocol", protocol)
  141. w.Header().Add("BITS-Session-Id", uuid)
  142. w.Header().Add("Accept-Encoding", "Identity")
  143. w.Write(nil)
  144. }
  145. func (b *BITSHandler) bitsFragment(w http.ResponseWriter, r *http.Request, uuid string) {
  146. var err error
  147. // Check for correct session
  148. if uuid == "" {
  149. bitsError(w, "", http.StatusBadRequest, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  150. return
  151. }
  152. // Check for existing session
  153. var srcDir string
  154. srcDir = path.Join(b.cfg.TempDir, uuid)
  155. if b, _ := exists(srcDir); !b {
  156. bitsError(w, uuid, http.StatusBadRequest, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  157. return
  158. }
  159. // Get filename and make sure the path is correct
  160. _, filename := path.Split(r.RequestURI)
  161. if filename == "" {
  162. bitsError(w, uuid, http.StatusBadRequest, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  163. return
  164. }
  165. // See if filename is blacklisted. If so, return an error
  166. for _, reg := range b.cfg.Disallowed {
  167. if b, _ := regexp.MatchString(reg, filename); b {
  168. // File is blacklisted
  169. bitsError(w, uuid, http.StatusBadRequest, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  170. return
  171. }
  172. }
  173. // See if filename is whitelisted
  174. allowed := false
  175. for _, reg := range b.cfg.Allowed {
  176. if b, _ := regexp.MatchString(reg, filename); b {
  177. allowed = true
  178. break
  179. }
  180. }
  181. if !allowed {
  182. // No whitelisting rules matched!
  183. bitsError(w, uuid, http.StatusBadRequest, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  184. return
  185. }
  186. // Get absolute paths to file
  187. var src string
  188. if src, err = filepath.Abs(filepath.Join(srcDir, filename)); err != nil {
  189. src = filepath.Join(srcDir, filename)
  190. }
  191. // Parse range
  192. var rangeStart, rangeEnd, fileLength uint64
  193. if rangeStart, rangeEnd, fileLength, err = parseRange(r.Header.Get("Content-Range")); err != nil {
  194. bitsError(w, uuid, http.StatusBadRequest, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  195. return
  196. }
  197. // Check filesize
  198. if fileLength > b.cfg.MaxSize && b.cfg.MaxSize > 0 {
  199. bitsError(w, uuid, http.StatusRequestEntityTooLarge, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  200. return
  201. }
  202. // Get the length of the posted data
  203. var fragmentSize uint64
  204. if fragmentSize, err = strconv.ParseUint(r.Header.Get("Content-Length"), 10, 64); err != nil {
  205. bitsError(w, uuid, http.StatusBadRequest, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  206. return
  207. }
  208. // Get posted data and confirm size
  209. data, err := ioutil.ReadAll(r.Body)
  210. if uint64(len(data)) != fragmentSize {
  211. bitsError(w, uuid, http.StatusBadRequest, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  212. return
  213. }
  214. // Check that content-range matches content-length
  215. if rangeEnd-rangeStart+1 != fragmentSize {
  216. bitsError(w, uuid, http.StatusBadRequest, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  217. return
  218. }
  219. // Open file
  220. var file *os.File
  221. var fileSize uint64
  222. if exist, _ := exists(src); !exist {
  223. // Create file
  224. if file, err = os.OpenFile(src, os.O_CREATE|os.O_WRONLY, 0600); err != nil {
  225. bitsError(w, uuid, http.StatusInternalServerError, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  226. return
  227. }
  228. // New file, size is zero
  229. fileSize = 0
  230. } else {
  231. // Open file for append
  232. if file, err = os.OpenFile(src, os.O_APPEND|os.O_WRONLY, 0666); err != nil {
  233. bitsError(w, uuid, http.StatusInternalServerError, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  234. return
  235. }
  236. // Get size on disk
  237. if info, err := file.Stat(); err != nil {
  238. file.Close()
  239. bitsError(w, uuid, http.StatusInternalServerError, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  240. return
  241. } else {
  242. fileSize = uint64(info.Size())
  243. }
  244. }
  245. defer file.Close()
  246. // Sanity checks
  247. if rangeEnd < fileSize {
  248. // The range is already written to disk
  249. w.Header().Add("BITS-Recieved-Content-Range", strconv.FormatUint(fileSize, 10))
  250. bitsError(w, uuid, http.StatusRequestedRangeNotSatisfiable, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  251. return
  252. } else if rangeStart > fileSize {
  253. // start must be <= fileSize, else there will be a gap
  254. w.Header().Add("BITS-Recieved-Content-Range", strconv.FormatUint(fileSize, 10))
  255. bitsError(w, uuid, http.StatusRequestedRangeNotSatisfiable, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  256. return
  257. }
  258. // Calculate the offset in the slice, if overlapping
  259. var dataOffset = fileSize - rangeStart
  260. // Write the data to file
  261. var written uint64
  262. if wr, err := file.Write(data[dataOffset:]); err != nil {
  263. bitsError(w, uuid, http.StatusInternalServerError, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  264. return
  265. } else {
  266. written = uint64(wr)
  267. }
  268. // Make sure we wrote everything we wanted
  269. if written != fragmentSize-dataOffset {
  270. bitsError(w, uuid, http.StatusInternalServerError, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  271. return
  272. }
  273. // Check if we have written everything
  274. if rangeEnd+1 == fileLength {
  275. // File is done! Move it!
  276. file.Close()
  277. // Call the callback
  278. b.callback(EventRecieveFile, uuid, src)
  279. }
  280. w.Header().Add("BITS-Packet-Type", "Ack")
  281. w.Header().Add("BITS-Session-Id", uuid)
  282. w.Header().Add("BITS-Received-Content-Range", strconv.FormatUint(fileSize+uint64(written), 10))
  283. w.Write(nil)
  284. }
  285. func (b *BITSHandler) bitsCancel(w http.ResponseWriter, r *http.Request, uuid string) {
  286. // Check for correct session
  287. if uuid == "" {
  288. bitsError(w, uuid, http.StatusBadRequest, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  289. return
  290. }
  291. destDir := path.Join(b.cfg.TempDir, uuid)
  292. if exist, _ := exists(destDir); !exist {
  293. bitsError(w, uuid, http.StatusBadRequest, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  294. return
  295. }
  296. b.callback(EventCancelSession, uuid, destDir)
  297. w.Header().Add("BITS-Packet-Type", "Ack")
  298. w.Header().Add("BITS-Session-Id", uuid)
  299. w.Write(nil)
  300. }
  301. func (b *BITSHandler) bitsClose(w http.ResponseWriter, r *http.Request, uuid string) {
  302. // Check for correct session
  303. if uuid == "" {
  304. bitsError(w, uuid, http.StatusBadRequest, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  305. return
  306. }
  307. destDir := path.Join(b.cfg.TempDir, uuid)
  308. if exist, _ := exists(destDir); !exist {
  309. bitsError(w, uuid, http.StatusBadRequest, 0, BG_ERROR_CONTEXT_REMOTE_FILE)
  310. return
  311. }
  312. b.callback(EventCloseSession, uuid, destDir)
  313. w.Header().Add("BITS-Packet-Type", "Ack")
  314. w.Header().Add("BITS-Session-Id", uuid)
  315. w.Write(nil)
  316. }
  317. func newUUID() (string, error) {
  318. // Stolen from http://play.golang.org/p/4FkNSiUDMg
  319. uuid := make([]byte, 16)
  320. if n, err := io.ReadFull(rand.Reader, uuid); n != len(uuid) || err != nil {
  321. return "", err
  322. }
  323. // https://tools.ietf.org/html/rfc4122#section-4.1.1
  324. uuid[8] = uuid[8]&^0xc0 | 0x80
  325. // https://tools.ietf.org/html/rfc4122#section-4.1.3
  326. uuid[6] = uuid[6]&^0xf0 | 0x40
  327. return fmt.Sprintf("%x-%x-%x-%x-%x", uuid[0:4], uuid[4:6], uuid[6:8], uuid[8:10], uuid[10:]), nil
  328. }
  329. func exists(path string) (bool, error) {
  330. var err error
  331. if _, err = os.Stat(path); err != nil && os.IsNotExist(err) {
  332. return false, nil
  333. }
  334. return true, err
  335. }
  336. func parseRange(rangeString string) (rangeStart, rangeEnd, fileLength uint64, err error) {
  337. // We only support "range #-#/#" syntax
  338. if !strings.HasPrefix(rangeString, "bytes ") {
  339. return 0, 0, 0, errors.New("invalid range syntax")
  340. }
  341. // Remove leading 6 characters
  342. rangeArray := strings.Split(rangeString[6:], "/")
  343. if len(rangeArray) != 2 {
  344. return 0, 0, 0, errors.New("invalid range syntax")
  345. }
  346. // Parse total length
  347. if fileLength, err = strconv.ParseUint(rangeArray[1], 10, 64); err != nil {
  348. return 0, 0, 0, err
  349. }
  350. // Get start and end of range
  351. rangeArray = strings.Split(rangeArray[0], "-")
  352. if len(rangeArray) != 2 {
  353. return 0, 0, 0, errors.New("invalid range syntax")
  354. }
  355. // Parse start value
  356. if rangeStart, err = strconv.ParseUint(rangeArray[0], 10, 64); err != nil {
  357. return 0, 0, 0, err
  358. }
  359. // Parse end value
  360. if rangeEnd, err = strconv.ParseUint(rangeArray[1], 10, 64); err != nil {
  361. return 0, 0, 0, err
  362. }
  363. // Return values
  364. return rangeStart, rangeEnd, fileLength, nil
  365. }