net.go 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899
  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. err = c.c.SetDeadline(time.Now().Add(c.timeout))
  48. if err != nil {
  49. return
  50. }
  51. return c.c.Read(b)
  52. }
  53. func (c *timeoutConn) Write(b []byte) (n int, err error) {
  54. if c.buf != nil {
  55. if !c.packetPending {
  56. c.buf.BeginPacket(packPrelogin)
  57. c.packetPending = true
  58. }
  59. n, err = c.buf.Write(b)
  60. if err != nil {
  61. return
  62. }
  63. return
  64. }
  65. err = c.c.SetDeadline(time.Now().Add(c.timeout))
  66. if err != nil {
  67. return
  68. }
  69. return c.c.Write(b)
  70. }
  71. func (c timeoutConn) Close() error {
  72. return c.c.Close()
  73. }
  74. func (c timeoutConn) LocalAddr() net.Addr {
  75. return c.c.LocalAddr()
  76. }
  77. func (c timeoutConn) RemoteAddr() net.Addr {
  78. return c.c.RemoteAddr()
  79. }
  80. func (c timeoutConn) SetDeadline(t time.Time) error {
  81. panic("Not implemented")
  82. }
  83. func (c timeoutConn) SetReadDeadline(t time.Time) error {
  84. panic("Not implemented")
  85. }
  86. func (c timeoutConn) SetWriteDeadline(t time.Time) error {
  87. panic("Not implemented")
  88. }