packet.go 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191
  1. package tftp
  2. import (
  3. "bytes"
  4. "encoding/binary"
  5. "fmt"
  6. )
  7. const (
  8. opRRQ = uint16(1) // Read request (RRQ)
  9. opWRQ = uint16(2) // Write request (WRQ)
  10. opDATA = uint16(3) // Data
  11. opACK = uint16(4) // Acknowledgement
  12. opERROR = uint16(5) // Error
  13. opOACK = uint16(6) // Options Acknowledgment
  14. )
  15. const (
  16. blockLength = 512
  17. datagramLength = 516
  18. )
  19. type options map[string]string
  20. // RRQ/WRQ packet
  21. //
  22. // 2 bytes string 1 byte string 1 byte
  23. // --------------------------------------------------
  24. // | Opcode | Filename | 0 | Mode | 0 |
  25. // --------------------------------------------------
  26. type pRRQ []byte
  27. type pWRQ []byte
  28. // packRQ returns length of the packet in b
  29. func packRQ(p []byte, op uint16, filename, mode string, opts options) int {
  30. binary.BigEndian.PutUint16(p, op)
  31. n := 2
  32. n += copy(p[2:len(p)-10], filename)
  33. p[n] = 0
  34. n++
  35. n += copy(p[n:], mode)
  36. p[n] = 0
  37. n++
  38. for name, value := range opts {
  39. n += copy(p[n:], name)
  40. p[n] = 0
  41. n++
  42. n += copy(p[n:], value)
  43. p[n] = 0
  44. n++
  45. }
  46. return n
  47. }
  48. func unpackRQ(p []byte) (filename, mode string, opts options, err error) {
  49. bs := bytes.Split(p[2:], []byte{0})
  50. if len(bs) < 2 {
  51. return "", "", nil, fmt.Errorf("missing filename or mode")
  52. }
  53. filename = string(bs[0])
  54. mode = string(bs[1])
  55. if len(bs) < 4 {
  56. return filename, mode, nil, nil
  57. }
  58. opts = make(options)
  59. for i := 2; i+1 < len(bs); i += 2 {
  60. opts[string(bs[i])] = string(bs[i+1])
  61. }
  62. return filename, mode, opts, nil
  63. }
  64. // OACK packet
  65. //
  66. // +----------+---~~---+---+---~~---+---+---~~---+---+---~~---+---+
  67. // | Opcode | opt1 | 0 | value1 | 0 | optN | 0 | valueN | 0 |
  68. // +----------+---~~---+---+---~~---+---+---~~---+---+---~~---+---+
  69. type pOACK []byte
  70. func packOACK(p []byte, opts options) int {
  71. binary.BigEndian.PutUint16(p, opOACK)
  72. n := 2
  73. for name, value := range opts {
  74. n += copy(p[n:], name)
  75. p[n] = 0
  76. n++
  77. n += copy(p[n:], value)
  78. p[n] = 0
  79. n++
  80. }
  81. return n
  82. }
  83. func unpackOACK(p []byte) (opts options, err error) {
  84. bs := bytes.Split(p[2:], []byte{0})
  85. opts = make(options)
  86. for i := 0; i+1 < len(bs); i += 2 {
  87. opts[string(bs[i])] = string(bs[i+1])
  88. }
  89. return opts, nil
  90. }
  91. // ERROR packet
  92. //
  93. // 2 bytes 2 bytes string 1 byte
  94. // ------------------------------------------
  95. // | Opcode | ErrorCode | ErrMsg | 0 |
  96. // ------------------------------------------
  97. type pERROR []byte
  98. func packERROR(p []byte, code uint16, message string) int {
  99. binary.BigEndian.PutUint16(p, opERROR)
  100. binary.BigEndian.PutUint16(p[2:], code)
  101. n := copy(p[4:len(p)-2], message)
  102. p[4+n] = 0
  103. return n + 5
  104. }
  105. func (p pERROR) code() uint16 {
  106. return binary.BigEndian.Uint16(p[2:])
  107. }
  108. func (p pERROR) message() string {
  109. return string(p[4:])
  110. }
  111. // DATA packet
  112. //
  113. // 2 bytes 2 bytes n bytes
  114. // ----------------------------------
  115. // | Opcode | Block # | Data |
  116. // ----------------------------------
  117. type pDATA []byte
  118. func (p pDATA) block() uint16 {
  119. return binary.BigEndian.Uint16(p[2:])
  120. }
  121. // ACK packet
  122. //
  123. // 2 bytes 2 bytes
  124. // -----------------------
  125. // | Opcode | Block # |
  126. // -----------------------
  127. type pACK []byte
  128. func (p pACK) block() uint16 {
  129. return binary.BigEndian.Uint16(p[2:])
  130. }
  131. func parsePacket(p []byte) (interface{}, error) {
  132. l := len(p)
  133. if l < 2 {
  134. return nil, fmt.Errorf("short packet")
  135. }
  136. opcode := binary.BigEndian.Uint16(p)
  137. switch opcode {
  138. case opRRQ:
  139. if l < 4 {
  140. return nil, fmt.Errorf("short RRQ packet: %d", l)
  141. }
  142. return pRRQ(p), nil
  143. case opWRQ:
  144. if l < 4 {
  145. return nil, fmt.Errorf("short WRQ packet: %d", l)
  146. }
  147. return pWRQ(p), nil
  148. case opDATA:
  149. if l < 4 {
  150. return nil, fmt.Errorf("short DATA packet: %d", l)
  151. }
  152. return pDATA(p), nil
  153. case opACK:
  154. if l < 4 {
  155. return nil, fmt.Errorf("short ACK packet: %d", l)
  156. }
  157. return pACK(p), nil
  158. case opERROR:
  159. if l < 5 {
  160. return nil, fmt.Errorf("short ERROR packet: %d", l)
  161. }
  162. return pERROR(p), nil
  163. case opOACK:
  164. if l < 6 {
  165. return nil, fmt.Errorf("short OACK packet: %d", l)
  166. }
  167. return pOACK(p), nil
  168. default:
  169. return nil, fmt.Errorf("unknown opcode: %d", opcode)
  170. }
  171. }