protocol.go 22 KB

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