protocol.go 23 KB

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