protocol.go 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940
  1. // Copyright (C) 2014 The Protocol Authors.
  2. package protocol
  3. import (
  4. "crypto/sha256"
  5. "encoding/binary"
  6. "errors"
  7. "fmt"
  8. "io"
  9. "path"
  10. "strings"
  11. "sync"
  12. "time"
  13. lz4 "github.com/bkaradzic/go-lz4"
  14. )
  15. const (
  16. // Shifts
  17. KiB = 10
  18. MiB = 20
  19. GiB = 30
  20. )
  21. const (
  22. // MaxMessageLen is the largest message size allowed on the wire. (500 MB)
  23. MaxMessageLen = 500 * 1000 * 1000
  24. // MinBlockSize is the minimum block size allowed
  25. MinBlockSize = 128 << KiB
  26. // MaxBlockSize is the maximum block size allowed
  27. MaxBlockSize = 16 << MiB
  28. // DesiredPerFileBlocks is the number of blocks we aim for per file
  29. DesiredPerFileBlocks = 2000
  30. )
  31. // BlockSizes is the list of valid block sizes, from min to max
  32. var BlockSizes []int
  33. // For each block size, the hash of a block of all zeroes
  34. var sha256OfEmptyBlock = make(map[int][sha256.Size]byte)
  35. func init() {
  36. for blockSize := MinBlockSize; blockSize <= MaxBlockSize; blockSize *= 2 {
  37. BlockSizes = append(BlockSizes, blockSize)
  38. sha256OfEmptyBlock[blockSize] = sha256.Sum256(make([]byte, blockSize))
  39. }
  40. BufferPool = newBufferPool()
  41. }
  42. // BlockSize returns the block size to use for the given file size
  43. func BlockSize(fileSize int64) int {
  44. var blockSize int
  45. for _, blockSize = range BlockSizes {
  46. if fileSize < DesiredPerFileBlocks*int64(blockSize) {
  47. break
  48. }
  49. }
  50. return blockSize
  51. }
  52. const (
  53. stateInitial = iota
  54. stateReady
  55. )
  56. // Request message flags
  57. const (
  58. FlagFromTemporary uint32 = 1 << iota
  59. )
  60. // ClusterConfigMessage.Folders flags
  61. const (
  62. FlagFolderReadOnly uint32 = 1 << 0
  63. FlagFolderIgnorePerms = 1 << 1
  64. FlagFolderIgnoreDelete = 1 << 2
  65. FlagFolderDisabledTempIndexes = 1 << 3
  66. FlagFolderAll = 1<<4 - 1
  67. )
  68. // ClusterConfigMessage.Folders.Devices flags
  69. const (
  70. FlagShareTrusted uint32 = 1 << 0
  71. FlagShareReadOnly = 1 << 1
  72. FlagIntroducer = 1 << 2
  73. FlagShareBits = 0x000000ff
  74. )
  75. // FileInfo.LocalFlags flags
  76. const (
  77. FlagLocalUnsupported = 1 << 0 // The kind is unsupported, e.g. symlinks on Windows
  78. FlagLocalIgnored = 1 << 1 // Matches local ignore patterns
  79. FlagLocalMustRescan = 1 << 2 // Doesn't match content on disk, must be rechecked fully
  80. FlagLocalReceiveOnly = 1 << 3 // Change detected on receive only folder
  81. // Flags that should result in the Invalid bit on outgoing updates
  82. LocalInvalidFlags = FlagLocalUnsupported | FlagLocalIgnored | FlagLocalMustRescan | FlagLocalReceiveOnly
  83. // Flags that should result in a file being in conflict with its
  84. // successor, due to us not having an up to date picture of its state on
  85. // disk.
  86. LocalConflictFlags = FlagLocalUnsupported | FlagLocalIgnored | FlagLocalReceiveOnly
  87. LocalAllFlags = FlagLocalUnsupported | FlagLocalIgnored | FlagLocalMustRescan | FlagLocalReceiveOnly
  88. )
  89. var (
  90. ErrClosed = errors.New("connection closed")
  91. ErrTimeout = errors.New("read timeout")
  92. ErrSwitchingConnections = errors.New("switching connections")
  93. errUnknownMessage = errors.New("unknown message")
  94. errInvalidFilename = errors.New("filename is invalid")
  95. errUncleanFilename = errors.New("filename not in canonical format")
  96. errDeletedHasBlocks = errors.New("deleted file with non-empty block list")
  97. errDirectoryHasBlocks = errors.New("directory with non-empty block list")
  98. errFileHasNoBlocks = errors.New("file with empty block list")
  99. )
  100. type Model interface {
  101. // An index was received from the peer device
  102. Index(deviceID DeviceID, folder string, files []FileInfo)
  103. // An index update was received from the peer device
  104. IndexUpdate(deviceID DeviceID, folder string, files []FileInfo)
  105. // A request was made by the peer device
  106. Request(deviceID DeviceID, folder, name string, size int32, offset int64, hash []byte, weakHash uint32, fromTemporary bool) (RequestResponse, error)
  107. // A cluster configuration message was received
  108. ClusterConfig(deviceID DeviceID, config ClusterConfig)
  109. // The peer device closed the connection
  110. Closed(conn Connection, err error)
  111. // The peer device sent progress updates for the files it is currently downloading
  112. DownloadProgress(deviceID DeviceID, folder string, updates []FileDownloadProgressUpdate)
  113. }
  114. type RequestResponse interface {
  115. Data() []byte
  116. Close() // Must always be called once the byte slice is no longer in use
  117. Wait() // Blocks until Close is called
  118. }
  119. type Connection interface {
  120. Start()
  121. Close(err error)
  122. ID() DeviceID
  123. Name() string
  124. Index(folder string, files []FileInfo) error
  125. IndexUpdate(folder string, files []FileInfo) error
  126. Request(folder string, name string, offset int64, size int, hash []byte, weakHash uint32, fromTemporary bool) ([]byte, error)
  127. ClusterConfig(config ClusterConfig)
  128. DownloadProgress(folder string, updates []FileDownloadProgressUpdate)
  129. Statistics() Statistics
  130. Closed() bool
  131. }
  132. type rawConnection struct {
  133. id DeviceID
  134. name string
  135. receiver Model
  136. cr *countingReader
  137. cw *countingWriter
  138. awaiting map[int32]chan asyncResult
  139. awaitingMut sync.Mutex
  140. idxMut sync.Mutex // ensures serialization of Index calls
  141. nextID int32
  142. nextIDMut sync.Mutex
  143. outbox chan asyncMessage
  144. sendClose chan asyncMessage
  145. closed chan struct{}
  146. once sync.Once
  147. compression Compression
  148. }
  149. type asyncResult struct {
  150. val []byte
  151. err error
  152. }
  153. type message interface {
  154. ProtoSize() int
  155. Marshal() ([]byte, error)
  156. MarshalTo([]byte) (int, error)
  157. Unmarshal([]byte) error
  158. }
  159. type asyncMessage struct {
  160. msg message
  161. done chan struct{} // done closes when we're done sending the message
  162. }
  163. const (
  164. // PingSendInterval is how often we make sure to send a message, by
  165. // triggering pings if necessary.
  166. PingSendInterval = 90 * time.Second
  167. // ReceiveTimeout is the longest we'll wait for a message from the other
  168. // side before closing the connection.
  169. ReceiveTimeout = 300 * time.Second
  170. )
  171. func NewConnection(deviceID DeviceID, reader io.Reader, writer io.Writer, receiver Model, name string, compress Compression) Connection {
  172. cr := &countingReader{Reader: reader}
  173. cw := &countingWriter{Writer: writer}
  174. c := rawConnection{
  175. id: deviceID,
  176. name: name,
  177. receiver: nativeModel{receiver},
  178. cr: cr,
  179. cw: cw,
  180. awaiting: make(map[int32]chan asyncResult),
  181. outbox: make(chan asyncMessage),
  182. sendClose: make(chan asyncMessage),
  183. closed: make(chan struct{}),
  184. compression: compress,
  185. }
  186. return wireFormatConnection{&c}
  187. }
  188. // Start creates the goroutines for sending and receiving of messages. It must
  189. // be called exactly once after creating a connection.
  190. func (c *rawConnection) Start() {
  191. go c.readerLoop()
  192. go c.writerLoop()
  193. go c.pingSender()
  194. go c.pingReceiver()
  195. }
  196. func (c *rawConnection) ID() DeviceID {
  197. return c.id
  198. }
  199. func (c *rawConnection) Name() string {
  200. return c.name
  201. }
  202. // Index writes the list of file information to the connected peer device
  203. func (c *rawConnection) Index(folder string, idx []FileInfo) error {
  204. select {
  205. case <-c.closed:
  206. return ErrClosed
  207. default:
  208. }
  209. c.idxMut.Lock()
  210. c.send(&Index{
  211. Folder: folder,
  212. Files: idx,
  213. }, nil)
  214. c.idxMut.Unlock()
  215. return nil
  216. }
  217. // IndexUpdate writes the list of file information to the connected peer device as an update
  218. func (c *rawConnection) IndexUpdate(folder string, idx []FileInfo) error {
  219. select {
  220. case <-c.closed:
  221. return ErrClosed
  222. default:
  223. }
  224. c.idxMut.Lock()
  225. c.send(&IndexUpdate{
  226. Folder: folder,
  227. Files: idx,
  228. }, nil)
  229. c.idxMut.Unlock()
  230. return nil
  231. }
  232. // Request returns the bytes for the specified block after fetching them from the connected peer.
  233. func (c *rawConnection) Request(folder string, name string, offset int64, size int, hash []byte, weakHash uint32, fromTemporary bool) ([]byte, error) {
  234. c.nextIDMut.Lock()
  235. id := c.nextID
  236. c.nextID++
  237. c.nextIDMut.Unlock()
  238. c.awaitingMut.Lock()
  239. if _, ok := c.awaiting[id]; ok {
  240. panic("id taken")
  241. }
  242. rc := make(chan asyncResult, 1)
  243. c.awaiting[id] = rc
  244. c.awaitingMut.Unlock()
  245. ok := c.send(&Request{
  246. ID: id,
  247. Folder: folder,
  248. Name: name,
  249. Offset: offset,
  250. Size: int32(size),
  251. Hash: hash,
  252. WeakHash: weakHash,
  253. FromTemporary: fromTemporary,
  254. }, nil)
  255. if !ok {
  256. return nil, ErrClosed
  257. }
  258. res, ok := <-rc
  259. if !ok {
  260. return nil, ErrClosed
  261. }
  262. return res.val, res.err
  263. }
  264. // ClusterConfig send the cluster configuration message to the peer and returns any error
  265. func (c *rawConnection) ClusterConfig(config ClusterConfig) {
  266. c.send(&config, nil)
  267. }
  268. func (c *rawConnection) Closed() bool {
  269. select {
  270. case <-c.closed:
  271. return true
  272. default:
  273. return false
  274. }
  275. }
  276. // DownloadProgress sends the progress updates for the files that are currently being downloaded.
  277. func (c *rawConnection) DownloadProgress(folder string, updates []FileDownloadProgressUpdate) {
  278. c.send(&DownloadProgress{
  279. Folder: folder,
  280. Updates: updates,
  281. }, nil)
  282. }
  283. func (c *rawConnection) ping() bool {
  284. return c.send(&Ping{}, nil)
  285. }
  286. func (c *rawConnection) readerLoop() (err error) {
  287. defer func() {
  288. c.internalClose(err)
  289. }()
  290. fourByteBuf := make([]byte, 4)
  291. state := stateInitial
  292. for {
  293. select {
  294. case <-c.closed:
  295. return ErrClosed
  296. default:
  297. }
  298. msg, err := c.readMessage(fourByteBuf)
  299. if err == errUnknownMessage {
  300. // Unknown message types are skipped, for future extensibility.
  301. continue
  302. }
  303. if err != nil {
  304. return err
  305. }
  306. switch msg := msg.(type) {
  307. case *ClusterConfig:
  308. l.Debugln("read ClusterConfig message")
  309. if state != stateInitial {
  310. return fmt.Errorf("protocol error: cluster config message in state %d", state)
  311. }
  312. c.receiver.ClusterConfig(c.id, *msg)
  313. state = stateReady
  314. case *Index:
  315. l.Debugln("read Index message")
  316. if state != stateReady {
  317. return fmt.Errorf("protocol error: index message in state %d", state)
  318. }
  319. if err := checkIndexConsistency(msg.Files); err != nil {
  320. return fmt.Errorf("protocol error: index: %v", err)
  321. }
  322. c.handleIndex(*msg)
  323. state = stateReady
  324. case *IndexUpdate:
  325. l.Debugln("read IndexUpdate message")
  326. if state != stateReady {
  327. return fmt.Errorf("protocol error: index update message in state %d", state)
  328. }
  329. if err := checkIndexConsistency(msg.Files); err != nil {
  330. return fmt.Errorf("protocol error: index update: %v", err)
  331. }
  332. c.handleIndexUpdate(*msg)
  333. state = stateReady
  334. case *Request:
  335. l.Debugln("read Request message")
  336. if state != stateReady {
  337. return fmt.Errorf("protocol error: request message in state %d", state)
  338. }
  339. if err := checkFilename(msg.Name); err != nil {
  340. return fmt.Errorf("protocol error: request: %q: %v", msg.Name, err)
  341. }
  342. go c.handleRequest(*msg)
  343. case *Response:
  344. l.Debugln("read Response message")
  345. if state != stateReady {
  346. return fmt.Errorf("protocol error: response message in state %d", state)
  347. }
  348. c.handleResponse(*msg)
  349. case *DownloadProgress:
  350. l.Debugln("read DownloadProgress message")
  351. if state != stateReady {
  352. return fmt.Errorf("protocol error: response message in state %d", state)
  353. }
  354. c.receiver.DownloadProgress(c.id, msg.Folder, msg.Updates)
  355. case *Ping:
  356. l.Debugln("read Ping message")
  357. if state != stateReady {
  358. return fmt.Errorf("protocol error: ping message in state %d", state)
  359. }
  360. // Nothing
  361. case *Close:
  362. l.Debugln("read Close message")
  363. return errors.New(msg.Reason)
  364. default:
  365. l.Debugf("read unknown message: %+T", msg)
  366. return fmt.Errorf("protocol error: %s: unknown or empty message", c.id)
  367. }
  368. }
  369. }
  370. func (c *rawConnection) readMessage(fourByteBuf []byte) (message, error) {
  371. hdr, err := c.readHeader(fourByteBuf)
  372. if err != nil {
  373. return nil, err
  374. }
  375. return c.readMessageAfterHeader(hdr, fourByteBuf)
  376. }
  377. func (c *rawConnection) readMessageAfterHeader(hdr Header, fourByteBuf []byte) (message, error) {
  378. // First comes a 4 byte message length
  379. if _, err := io.ReadFull(c.cr, fourByteBuf[:4]); err != nil {
  380. return nil, fmt.Errorf("reading message length: %v", err)
  381. }
  382. msgLen := int32(binary.BigEndian.Uint32(fourByteBuf))
  383. if msgLen < 0 {
  384. return nil, fmt.Errorf("negative message length %d", msgLen)
  385. }
  386. // Then comes the message
  387. buf := BufferPool.Get(int(msgLen))
  388. if _, err := io.ReadFull(c.cr, buf); err != nil {
  389. return nil, fmt.Errorf("reading message: %v", err)
  390. }
  391. // ... which might be compressed
  392. switch hdr.Compression {
  393. case MessageCompressionNone:
  394. // Nothing
  395. case MessageCompressionLZ4:
  396. decomp, err := c.lz4Decompress(buf)
  397. BufferPool.Put(buf)
  398. if err != nil {
  399. return nil, fmt.Errorf("decompressing message: %v", err)
  400. }
  401. buf = decomp
  402. default:
  403. return nil, fmt.Errorf("unknown message compression %d", hdr.Compression)
  404. }
  405. // ... and is then unmarshalled
  406. msg, err := c.newMessage(hdr.Type)
  407. if err != nil {
  408. return nil, err
  409. }
  410. if err := msg.Unmarshal(buf); err != nil {
  411. return nil, fmt.Errorf("unmarshalling message: %v", err)
  412. }
  413. BufferPool.Put(buf)
  414. return msg, nil
  415. }
  416. func (c *rawConnection) readHeader(fourByteBuf []byte) (Header, error) {
  417. // First comes a 2 byte header length
  418. if _, err := io.ReadFull(c.cr, fourByteBuf[:2]); err != nil {
  419. return Header{}, fmt.Errorf("reading length: %v", err)
  420. }
  421. hdrLen := int16(binary.BigEndian.Uint16(fourByteBuf))
  422. if hdrLen < 0 {
  423. return Header{}, fmt.Errorf("negative header length %d", hdrLen)
  424. }
  425. // Then comes the header
  426. buf := BufferPool.Get(int(hdrLen))
  427. if _, err := io.ReadFull(c.cr, buf); err != nil {
  428. return Header{}, fmt.Errorf("reading header: %v", err)
  429. }
  430. var hdr Header
  431. if err := hdr.Unmarshal(buf); err != nil {
  432. return Header{}, fmt.Errorf("unmarshalling header: %v", err)
  433. }
  434. BufferPool.Put(buf)
  435. return hdr, nil
  436. }
  437. func (c *rawConnection) handleIndex(im Index) {
  438. l.Debugf("Index(%v, %v, %d file)", c.id, im.Folder, len(im.Files))
  439. c.receiver.Index(c.id, im.Folder, im.Files)
  440. }
  441. func (c *rawConnection) handleIndexUpdate(im IndexUpdate) {
  442. l.Debugf("queueing IndexUpdate(%v, %v, %d files)", c.id, im.Folder, len(im.Files))
  443. c.receiver.IndexUpdate(c.id, im.Folder, im.Files)
  444. }
  445. // checkIndexConsistency verifies a number of invariants on FileInfos received in
  446. // index messages.
  447. func checkIndexConsistency(fs []FileInfo) error {
  448. for _, f := range fs {
  449. if err := checkFileInfoConsistency(f); err != nil {
  450. return fmt.Errorf("%q: %v", f.Name, err)
  451. }
  452. }
  453. return nil
  454. }
  455. // checkFileInfoConsistency verifies a number of invariants on the given FileInfo
  456. func checkFileInfoConsistency(f FileInfo) error {
  457. if err := checkFilename(f.Name); err != nil {
  458. return err
  459. }
  460. switch {
  461. case f.Deleted && len(f.Blocks) != 0:
  462. // Deleted files should have no blocks
  463. return errDeletedHasBlocks
  464. case f.Type == FileInfoTypeDirectory && len(f.Blocks) != 0:
  465. // Directories should have no blocks
  466. return errDirectoryHasBlocks
  467. case !f.Deleted && !f.IsInvalid() && f.Type == FileInfoTypeFile && len(f.Blocks) == 0:
  468. // Non-deleted, non-invalid files should have at least one block
  469. return errFileHasNoBlocks
  470. }
  471. return nil
  472. }
  473. // checkFilename verifies that the given filename is valid according to the
  474. // spec on what's allowed over the wire. A filename failing this test is
  475. // grounds for disconnecting the device.
  476. func checkFilename(name string) error {
  477. cleanedName := path.Clean(name)
  478. if cleanedName != name {
  479. // The filename on the wire should be in canonical format. If
  480. // Clean() managed to clean it up, there was something wrong with
  481. // it.
  482. return errUncleanFilename
  483. }
  484. switch name {
  485. case "", ".", "..":
  486. // These names are always invalid.
  487. return errInvalidFilename
  488. }
  489. if strings.HasPrefix(name, "/") {
  490. // Names are folder relative, not absolute.
  491. return errInvalidFilename
  492. }
  493. if strings.HasPrefix(name, "../") {
  494. // Starting with a dotdot is not allowed. Any other dotdots would
  495. // have been handled by the Clean() call at the top.
  496. return errInvalidFilename
  497. }
  498. return nil
  499. }
  500. func (c *rawConnection) handleRequest(req Request) {
  501. res, err := c.receiver.Request(c.id, req.Folder, req.Name, req.Size, req.Offset, req.Hash, req.WeakHash, req.FromTemporary)
  502. if err != nil {
  503. c.send(&Response{
  504. ID: req.ID,
  505. Code: errorToCode(err),
  506. }, nil)
  507. return
  508. }
  509. done := make(chan struct{})
  510. c.send(&Response{
  511. ID: req.ID,
  512. Data: res.Data(),
  513. Code: errorToCode(nil),
  514. }, done)
  515. <-done
  516. res.Close()
  517. }
  518. func (c *rawConnection) handleResponse(resp Response) {
  519. c.awaitingMut.Lock()
  520. if rc := c.awaiting[resp.ID]; rc != nil {
  521. delete(c.awaiting, resp.ID)
  522. rc <- asyncResult{resp.Data, codeToError(resp.Code)}
  523. close(rc)
  524. }
  525. c.awaitingMut.Unlock()
  526. }
  527. func (c *rawConnection) send(msg message, done chan struct{}) bool {
  528. select {
  529. case c.outbox <- asyncMessage{msg, done}:
  530. return true
  531. case <-c.closed:
  532. if done != nil {
  533. close(done)
  534. }
  535. return false
  536. }
  537. }
  538. func (c *rawConnection) writerLoop() {
  539. for {
  540. select {
  541. case hm := <-c.outbox:
  542. err := c.writeMessage(hm)
  543. if hm.done != nil {
  544. close(hm.done)
  545. }
  546. if err != nil {
  547. c.internalClose(err)
  548. return
  549. }
  550. case m := <-c.sendClose:
  551. c.writeMessage(m)
  552. close(m.done)
  553. return // No message must be sent after the Close message.
  554. case <-c.closed:
  555. return
  556. }
  557. }
  558. }
  559. func (c *rawConnection) writeMessage(hm asyncMessage) error {
  560. if c.shouldCompressMessage(hm.msg) {
  561. return c.writeCompressedMessage(hm)
  562. }
  563. return c.writeUncompressedMessage(hm)
  564. }
  565. func (c *rawConnection) writeCompressedMessage(hm asyncMessage) error {
  566. size := hm.msg.ProtoSize()
  567. buf := BufferPool.Get(size)
  568. if _, err := hm.msg.MarshalTo(buf); err != nil {
  569. return fmt.Errorf("marshalling message: %v", err)
  570. }
  571. compressed, err := c.lz4Compress(buf)
  572. if err != nil {
  573. return fmt.Errorf("compressing message: %v", err)
  574. }
  575. hdr := Header{
  576. Type: c.typeOf(hm.msg),
  577. Compression: MessageCompressionLZ4,
  578. }
  579. hdrSize := hdr.ProtoSize()
  580. if hdrSize > 1<<16-1 {
  581. panic("impossibly large header")
  582. }
  583. totSize := 2 + hdrSize + 4 + len(compressed)
  584. buf = BufferPool.Upgrade(buf, totSize)
  585. // Header length
  586. binary.BigEndian.PutUint16(buf, uint16(hdrSize))
  587. // Header
  588. if _, err := hdr.MarshalTo(buf[2:]); err != nil {
  589. return fmt.Errorf("marshalling header: %v", err)
  590. }
  591. // Message length
  592. binary.BigEndian.PutUint32(buf[2+hdrSize:], uint32(len(compressed)))
  593. // Message
  594. copy(buf[2+hdrSize+4:], compressed)
  595. BufferPool.Put(compressed)
  596. n, err := c.cw.Write(buf)
  597. BufferPool.Put(buf)
  598. l.Debugf("wrote %d bytes on the wire (2 bytes length, %d bytes header, 4 bytes message length, %d bytes message (%d uncompressed)), err=%v", n, hdrSize, len(compressed), size, err)
  599. if err != nil {
  600. return fmt.Errorf("writing message: %v", err)
  601. }
  602. return nil
  603. }
  604. func (c *rawConnection) writeUncompressedMessage(hm asyncMessage) error {
  605. size := hm.msg.ProtoSize()
  606. hdr := Header{
  607. Type: c.typeOf(hm.msg),
  608. }
  609. hdrSize := hdr.ProtoSize()
  610. if hdrSize > 1<<16-1 {
  611. panic("impossibly large header")
  612. }
  613. totSize := 2 + hdrSize + 4 + size
  614. buf := BufferPool.Get(totSize)
  615. // Header length
  616. binary.BigEndian.PutUint16(buf, uint16(hdrSize))
  617. // Header
  618. if _, err := hdr.MarshalTo(buf[2:]); err != nil {
  619. return fmt.Errorf("marshalling header: %v", err)
  620. }
  621. // Message length
  622. binary.BigEndian.PutUint32(buf[2+hdrSize:], uint32(size))
  623. // Message
  624. if _, err := hm.msg.MarshalTo(buf[2+hdrSize+4:]); err != nil {
  625. return fmt.Errorf("marshalling message: %v", err)
  626. }
  627. n, err := c.cw.Write(buf[:totSize])
  628. BufferPool.Put(buf)
  629. l.Debugf("wrote %d bytes on the wire (2 bytes length, %d bytes header, 4 bytes message length, %d bytes message), err=%v", n, hdrSize, size, err)
  630. if err != nil {
  631. return fmt.Errorf("writing message: %v", err)
  632. }
  633. return nil
  634. }
  635. func (c *rawConnection) typeOf(msg message) MessageType {
  636. switch msg.(type) {
  637. case *ClusterConfig:
  638. return messageTypeClusterConfig
  639. case *Index:
  640. return messageTypeIndex
  641. case *IndexUpdate:
  642. return messageTypeIndexUpdate
  643. case *Request:
  644. return messageTypeRequest
  645. case *Response:
  646. return messageTypeResponse
  647. case *DownloadProgress:
  648. return messageTypeDownloadProgress
  649. case *Ping:
  650. return messageTypePing
  651. case *Close:
  652. return messageTypeClose
  653. default:
  654. panic("bug: unknown message type")
  655. }
  656. }
  657. func (c *rawConnection) newMessage(t MessageType) (message, error) {
  658. switch t {
  659. case messageTypeClusterConfig:
  660. return new(ClusterConfig), nil
  661. case messageTypeIndex:
  662. return new(Index), nil
  663. case messageTypeIndexUpdate:
  664. return new(IndexUpdate), nil
  665. case messageTypeRequest:
  666. return new(Request), nil
  667. case messageTypeResponse:
  668. return new(Response), nil
  669. case messageTypeDownloadProgress:
  670. return new(DownloadProgress), nil
  671. case messageTypePing:
  672. return new(Ping), nil
  673. case messageTypeClose:
  674. return new(Close), nil
  675. default:
  676. return nil, errUnknownMessage
  677. }
  678. }
  679. func (c *rawConnection) shouldCompressMessage(msg message) bool {
  680. switch c.compression {
  681. case CompressNever:
  682. return false
  683. case CompressAlways:
  684. // Use compression for large enough messages
  685. return msg.ProtoSize() >= compressionThreshold
  686. case CompressMetadata:
  687. _, isResponse := msg.(*Response)
  688. // Compress if it's large enough and not a response message
  689. return !isResponse && msg.ProtoSize() >= compressionThreshold
  690. default:
  691. panic("unknown compression setting")
  692. }
  693. }
  694. // Close is called when the connection is regularely closed and thus the Close
  695. // BEP message is sent before terminating the actual connection. The error
  696. // argument specifies the reason for closing the connection.
  697. func (c *rawConnection) Close(err error) {
  698. c.once.Do(func() {
  699. done := make(chan struct{})
  700. c.sendClose <- asyncMessage{&Close{err.Error()}, done}
  701. <-done
  702. // No more sends are necessary, therefore closing the underlying
  703. // connection can happen at the same time as the internal cleanup.
  704. // And this prevents a potential deadlock due to calling c.receiver.Closed
  705. go c.commonClose(err)
  706. })
  707. }
  708. // internalClose is called if there is an unexpected error during normal operation.
  709. func (c *rawConnection) internalClose(err error) {
  710. c.once.Do(func() {
  711. c.commonClose(err)
  712. })
  713. }
  714. // commonClose is a utility function that must only be called from within
  715. // rawConnection.once.Do (i.e. in Close and close).
  716. func (c *rawConnection) commonClose(err error) {
  717. l.Debugln("close due to", err)
  718. close(c.closed)
  719. c.awaitingMut.Lock()
  720. for i, ch := range c.awaiting {
  721. if ch != nil {
  722. close(ch)
  723. delete(c.awaiting, i)
  724. }
  725. }
  726. c.awaitingMut.Unlock()
  727. c.receiver.Closed(c, err)
  728. }
  729. // The pingSender makes sure that we've sent a message within the last
  730. // PingSendInterval. If we already have something sent in the last
  731. // PingSendInterval/2, we do nothing. Otherwise we send a ping message. This
  732. // results in an effecting ping interval of somewhere between
  733. // PingSendInterval/2 and PingSendInterval.
  734. func (c *rawConnection) pingSender() {
  735. ticker := time.NewTicker(PingSendInterval / 2)
  736. defer ticker.Stop()
  737. for {
  738. select {
  739. case <-ticker.C:
  740. d := time.Since(c.cw.Last())
  741. if d < PingSendInterval/2 {
  742. l.Debugln(c.id, "ping skipped after wr", d)
  743. continue
  744. }
  745. l.Debugln(c.id, "ping -> after", d)
  746. c.ping()
  747. case <-c.closed:
  748. return
  749. }
  750. }
  751. }
  752. // The pingReceiver checks that we've received a message (any message will do,
  753. // but we expect pings in the absence of other messages) within the last
  754. // ReceiveTimeout. If not, we close the connection with an ErrTimeout.
  755. func (c *rawConnection) pingReceiver() {
  756. ticker := time.NewTicker(ReceiveTimeout / 2)
  757. defer ticker.Stop()
  758. for {
  759. select {
  760. case <-ticker.C:
  761. d := time.Since(c.cr.Last())
  762. if d > ReceiveTimeout {
  763. l.Debugln(c.id, "ping timeout", d)
  764. c.internalClose(ErrTimeout)
  765. }
  766. l.Debugln(c.id, "last read within", d)
  767. case <-c.closed:
  768. return
  769. }
  770. }
  771. }
  772. type Statistics struct {
  773. At time.Time
  774. InBytesTotal int64
  775. OutBytesTotal int64
  776. }
  777. func (c *rawConnection) Statistics() Statistics {
  778. return Statistics{
  779. At: time.Now(),
  780. InBytesTotal: c.cr.Tot(),
  781. OutBytesTotal: c.cw.Tot(),
  782. }
  783. }
  784. func (c *rawConnection) lz4Compress(src []byte) ([]byte, error) {
  785. var err error
  786. buf := BufferPool.Get(len(src))
  787. buf, err = lz4.Encode(buf, src)
  788. if err != nil {
  789. return nil, err
  790. }
  791. binary.BigEndian.PutUint32(buf, binary.LittleEndian.Uint32(buf))
  792. return buf, nil
  793. }
  794. func (c *rawConnection) lz4Decompress(src []byte) ([]byte, error) {
  795. size := binary.BigEndian.Uint32(src)
  796. binary.LittleEndian.PutUint32(src, size)
  797. var err error
  798. buf := BufferPool.Get(int(size))
  799. buf, err = lz4.Decode(buf, src)
  800. if err != nil {
  801. return nil, err
  802. }
  803. return buf, nil
  804. }