protocol.go 21 KB

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