net.go 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103
  1. package mssql
  2. import (
  3. "fmt"
  4. "net"
  5. "time"
  6. )
  7. type timeoutConn struct {
  8. c net.Conn
  9. timeout time.Duration
  10. buf *tdsBuffer
  11. packetPending bool
  12. continueRead bool
  13. }
  14. func newTimeoutConn(conn net.Conn, timeout time.Duration) *timeoutConn {
  15. return &timeoutConn{
  16. c: conn,
  17. timeout: timeout,
  18. }
  19. }
  20. func (c *timeoutConn) Read(b []byte) (n int, err error) {
  21. if c.buf != nil {
  22. if c.packetPending {
  23. c.packetPending = false
  24. err = c.buf.FinishPacket()
  25. if err != nil {
  26. err = fmt.Errorf("Cannot send handshake packet: %s", err.Error())
  27. return
  28. }
  29. c.continueRead = false
  30. }
  31. if !c.continueRead {
  32. var packet packetType
  33. packet, err = c.buf.BeginRead()
  34. if err != nil {
  35. err = fmt.Errorf("Cannot read handshake packet: %s", err.Error())
  36. return
  37. }
  38. if packet != packPrelogin {
  39. err = fmt.Errorf("unexpected packet %d, expecting prelogin", packet)
  40. return
  41. }
  42. c.continueRead = true
  43. }
  44. n, err = c.buf.Read(b)
  45. return
  46. }
  47. if c.timeout > 0 {
  48. err = c.c.SetDeadline(time.Now().Add(c.timeout))
  49. if err != nil {
  50. return
  51. }
  52. }
  53. return c.c.Read(b)
  54. }
  55. func (c *timeoutConn) Write(b []byte) (n int, err error) {
  56. if c.buf != nil {
  57. if !c.packetPending {
  58. c.buf.BeginPacket(packPrelogin, false)
  59. c.packetPending = true
  60. }
  61. n, err = c.buf.Write(b)
  62. if err != nil {
  63. return
  64. }
  65. return
  66. }
  67. if c.timeout > 0 {
  68. err = c.c.SetDeadline(time.Now().Add(c.timeout))
  69. if err != nil {
  70. return
  71. }
  72. }
  73. return c.c.Write(b)
  74. }
  75. func (c timeoutConn) Close() error {
  76. return c.c.Close()
  77. }
  78. func (c timeoutConn) LocalAddr() net.Addr {
  79. return c.c.LocalAddr()
  80. }
  81. func (c timeoutConn) RemoteAddr() net.Addr {
  82. return c.c.RemoteAddr()
  83. }
  84. func (c timeoutConn) SetDeadline(t time.Time) error {
  85. panic("Not implemented")
  86. }
  87. func (c timeoutConn) SetReadDeadline(t time.Time) error {
  88. panic("Not implemented")
  89. }
  90. func (c timeoutConn) SetWriteDeadline(t time.Time) error {
  91. panic("Not implemented")
  92. }