github.com/xushiwei/go@v0.0.0-20130601165731-2b9d83f45bc9/src/pkg/net/http/httputil/persist.go (about)

     1  // Copyright 2009 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package httputil provides HTTP utility functions, complementing the
     6  // more common ones in the net/http package.
     7  package httputil
     8  
     9  import (
    10  	"bufio"
    11  	"errors"
    12  	"io"
    13  	"net"
    14  	"net/http"
    15  	"net/textproto"
    16  	"sync"
    17  )
    18  
    19  var (
    20  	ErrPersistEOF = &http.ProtocolError{ErrorString: "persistent connection closed"}
    21  	ErrClosed     = &http.ProtocolError{ErrorString: "connection closed by user"}
    22  	ErrPipeline   = &http.ProtocolError{ErrorString: "pipeline error"}
    23  )
    24  
    25  // This is an API usage error - the local side is closed.
    26  // ErrPersistEOF (above) reports that the remote side is closed.
    27  var errClosed = errors.New("i/o operation on closed connection")
    28  
    29  // A ServerConn reads requests and sends responses over an underlying
    30  // connection, until the HTTP keepalive logic commands an end. ServerConn
    31  // also allows hijacking the underlying connection by calling Hijack
    32  // to regain control over the connection. ServerConn supports pipe-lining,
    33  // i.e. requests can be read out of sync (but in the same order) while the
    34  // respective responses are sent.
    35  //
    36  // ServerConn is low-level and should not be needed by most applications.
    37  // See Server.
    38  type ServerConn struct {
    39  	lk              sync.Mutex // read-write protects the following fields
    40  	c               net.Conn
    41  	r               *bufio.Reader
    42  	re, we          error // read/write errors
    43  	lastbody        io.ReadCloser
    44  	nread, nwritten int
    45  	pipereq         map[*http.Request]uint
    46  
    47  	pipe textproto.Pipeline
    48  }
    49  
    50  // NewServerConn returns a new ServerConn reading and writing c.  If r is not
    51  // nil, it is the buffer to use when reading c.
    52  func NewServerConn(c net.Conn, r *bufio.Reader) *ServerConn {
    53  	if r == nil {
    54  		r = bufio.NewReader(c)
    55  	}
    56  	return &ServerConn{c: c, r: r, pipereq: make(map[*http.Request]uint)}
    57  }
    58  
    59  // Hijack detaches the ServerConn and returns the underlying connection as well
    60  // as the read-side bufio which may have some left over data. Hijack may be
    61  // called before Read has signaled the end of the keep-alive logic. The user
    62  // should not call Hijack while Read or Write is in progress.
    63  func (sc *ServerConn) Hijack() (c net.Conn, r *bufio.Reader) {
    64  	sc.lk.Lock()
    65  	defer sc.lk.Unlock()
    66  	c = sc.c
    67  	r = sc.r
    68  	sc.c = nil
    69  	sc.r = nil
    70  	return
    71  }
    72  
    73  // Close calls Hijack and then also closes the underlying connection
    74  func (sc *ServerConn) Close() error {
    75  	c, _ := sc.Hijack()
    76  	if c != nil {
    77  		return c.Close()
    78  	}
    79  	return nil
    80  }
    81  
    82  // Read returns the next request on the wire. An ErrPersistEOF is returned if
    83  // it is gracefully determined that there are no more requests (e.g. after the
    84  // first request on an HTTP/1.0 connection, or after a Connection:close on a
    85  // HTTP/1.1 connection).
    86  func (sc *ServerConn) Read() (req *http.Request, err error) {
    87  
    88  	// Ensure ordered execution of Reads and Writes
    89  	id := sc.pipe.Next()
    90  	sc.pipe.StartRequest(id)
    91  	defer func() {
    92  		sc.pipe.EndRequest(id)
    93  		if req == nil {
    94  			sc.pipe.StartResponse(id)
    95  			sc.pipe.EndResponse(id)
    96  		} else {
    97  			// Remember the pipeline id of this request
    98  			sc.lk.Lock()
    99  			sc.pipereq[req] = id
   100  			sc.lk.Unlock()
   101  		}
   102  	}()
   103  
   104  	sc.lk.Lock()
   105  	if sc.we != nil { // no point receiving if write-side broken or closed
   106  		defer sc.lk.Unlock()
   107  		return nil, sc.we
   108  	}
   109  	if sc.re != nil {
   110  		defer sc.lk.Unlock()
   111  		return nil, sc.re
   112  	}
   113  	if sc.r == nil { // connection closed by user in the meantime
   114  		defer sc.lk.Unlock()
   115  		return nil, errClosed
   116  	}
   117  	r := sc.r
   118  	lastbody := sc.lastbody
   119  	sc.lastbody = nil
   120  	sc.lk.Unlock()
   121  
   122  	// Make sure body is fully consumed, even if user does not call body.Close
   123  	if lastbody != nil {
   124  		// body.Close is assumed to be idempotent and multiple calls to
   125  		// it should return the error that its first invocation
   126  		// returned.
   127  		err = lastbody.Close()
   128  		if err != nil {
   129  			sc.lk.Lock()
   130  			defer sc.lk.Unlock()
   131  			sc.re = err
   132  			return nil, err
   133  		}
   134  	}
   135  
   136  	req, err = http.ReadRequest(r)
   137  	sc.lk.Lock()
   138  	defer sc.lk.Unlock()
   139  	if err != nil {
   140  		if err == io.ErrUnexpectedEOF {
   141  			// A close from the opposing client is treated as a
   142  			// graceful close, even if there was some unparse-able
   143  			// data before the close.
   144  			sc.re = ErrPersistEOF
   145  			return nil, sc.re
   146  		} else {
   147  			sc.re = err
   148  			return req, err
   149  		}
   150  	}
   151  	sc.lastbody = req.Body
   152  	sc.nread++
   153  	if req.Close {
   154  		sc.re = ErrPersistEOF
   155  		return req, sc.re
   156  	}
   157  	return req, err
   158  }
   159  
   160  // Pending returns the number of unanswered requests
   161  // that have been received on the connection.
   162  func (sc *ServerConn) Pending() int {
   163  	sc.lk.Lock()
   164  	defer sc.lk.Unlock()
   165  	return sc.nread - sc.nwritten
   166  }
   167  
   168  // Write writes resp in response to req. To close the connection gracefully, set the
   169  // Response.Close field to true. Write should be considered operational until
   170  // it returns an error, regardless of any errors returned on the Read side.
   171  func (sc *ServerConn) Write(req *http.Request, resp *http.Response) error {
   172  
   173  	// Retrieve the pipeline ID of this request/response pair
   174  	sc.lk.Lock()
   175  	id, ok := sc.pipereq[req]
   176  	delete(sc.pipereq, req)
   177  	if !ok {
   178  		sc.lk.Unlock()
   179  		return ErrPipeline
   180  	}
   181  	sc.lk.Unlock()
   182  
   183  	// Ensure pipeline order
   184  	sc.pipe.StartResponse(id)
   185  	defer sc.pipe.EndResponse(id)
   186  
   187  	sc.lk.Lock()
   188  	if sc.we != nil {
   189  		defer sc.lk.Unlock()
   190  		return sc.we
   191  	}
   192  	if sc.c == nil { // connection closed by user in the meantime
   193  		defer sc.lk.Unlock()
   194  		return ErrClosed
   195  	}
   196  	c := sc.c
   197  	if sc.nread <= sc.nwritten {
   198  		defer sc.lk.Unlock()
   199  		return errors.New("persist server pipe count")
   200  	}
   201  	if resp.Close {
   202  		// After signaling a keep-alive close, any pipelined unread
   203  		// requests will be lost. It is up to the user to drain them
   204  		// before signaling.
   205  		sc.re = ErrPersistEOF
   206  	}
   207  	sc.lk.Unlock()
   208  
   209  	err := resp.Write(c)
   210  	sc.lk.Lock()
   211  	defer sc.lk.Unlock()
   212  	if err != nil {
   213  		sc.we = err
   214  		return err
   215  	}
   216  	sc.nwritten++
   217  
   218  	return nil
   219  }
   220  
   221  // A ClientConn sends request and receives headers over an underlying
   222  // connection, while respecting the HTTP keepalive logic. ClientConn
   223  // supports hijacking the connection calling Hijack to
   224  // regain control of the underlying net.Conn and deal with it as desired.
   225  //
   226  // ClientConn is low-level and should not be needed by most applications.
   227  // See Client.
   228  type ClientConn struct {
   229  	lk              sync.Mutex // read-write protects the following fields
   230  	c               net.Conn
   231  	r               *bufio.Reader
   232  	re, we          error // read/write errors
   233  	lastbody        io.ReadCloser
   234  	nread, nwritten int
   235  	pipereq         map[*http.Request]uint
   236  
   237  	pipe     textproto.Pipeline
   238  	writeReq func(*http.Request, io.Writer) error
   239  }
   240  
   241  // NewClientConn returns a new ClientConn reading and writing c.  If r is not
   242  // nil, it is the buffer to use when reading c.
   243  func NewClientConn(c net.Conn, r *bufio.Reader) *ClientConn {
   244  	if r == nil {
   245  		r = bufio.NewReader(c)
   246  	}
   247  	return &ClientConn{
   248  		c:        c,
   249  		r:        r,
   250  		pipereq:  make(map[*http.Request]uint),
   251  		writeReq: (*http.Request).Write,
   252  	}
   253  }
   254  
   255  // NewProxyClientConn works like NewClientConn but writes Requests
   256  // using Request's WriteProxy method.
   257  func NewProxyClientConn(c net.Conn, r *bufio.Reader) *ClientConn {
   258  	cc := NewClientConn(c, r)
   259  	cc.writeReq = (*http.Request).WriteProxy
   260  	return cc
   261  }
   262  
   263  // Hijack detaches the ClientConn and returns the underlying connection as well
   264  // as the read-side bufio which may have some left over data. Hijack may be
   265  // called before the user or Read have signaled the end of the keep-alive
   266  // logic. The user should not call Hijack while Read or Write is in progress.
   267  func (cc *ClientConn) Hijack() (c net.Conn, r *bufio.Reader) {
   268  	cc.lk.Lock()
   269  	defer cc.lk.Unlock()
   270  	c = cc.c
   271  	r = cc.r
   272  	cc.c = nil
   273  	cc.r = nil
   274  	return
   275  }
   276  
   277  // Close calls Hijack and then also closes the underlying connection
   278  func (cc *ClientConn) Close() error {
   279  	c, _ := cc.Hijack()
   280  	if c != nil {
   281  		return c.Close()
   282  	}
   283  	return nil
   284  }
   285  
   286  // Write writes a request. An ErrPersistEOF error is returned if the connection
   287  // has been closed in an HTTP keepalive sense. If req.Close equals true, the
   288  // keepalive connection is logically closed after this request and the opposing
   289  // server is informed. An ErrUnexpectedEOF indicates the remote closed the
   290  // underlying TCP connection, which is usually considered as graceful close.
   291  func (cc *ClientConn) Write(req *http.Request) (err error) {
   292  
   293  	// Ensure ordered execution of Writes
   294  	id := cc.pipe.Next()
   295  	cc.pipe.StartRequest(id)
   296  	defer func() {
   297  		cc.pipe.EndRequest(id)
   298  		if err != nil {
   299  			cc.pipe.StartResponse(id)
   300  			cc.pipe.EndResponse(id)
   301  		} else {
   302  			// Remember the pipeline id of this request
   303  			cc.lk.Lock()
   304  			cc.pipereq[req] = id
   305  			cc.lk.Unlock()
   306  		}
   307  	}()
   308  
   309  	cc.lk.Lock()
   310  	if cc.re != nil { // no point sending if read-side closed or broken
   311  		defer cc.lk.Unlock()
   312  		return cc.re
   313  	}
   314  	if cc.we != nil {
   315  		defer cc.lk.Unlock()
   316  		return cc.we
   317  	}
   318  	if cc.c == nil { // connection closed by user in the meantime
   319  		defer cc.lk.Unlock()
   320  		return errClosed
   321  	}
   322  	c := cc.c
   323  	if req.Close {
   324  		// We write the EOF to the write-side error, because there
   325  		// still might be some pipelined reads
   326  		cc.we = ErrPersistEOF
   327  	}
   328  	cc.lk.Unlock()
   329  
   330  	err = cc.writeReq(req, c)
   331  	cc.lk.Lock()
   332  	defer cc.lk.Unlock()
   333  	if err != nil {
   334  		cc.we = err
   335  		return err
   336  	}
   337  	cc.nwritten++
   338  
   339  	return nil
   340  }
   341  
   342  // Pending returns the number of unanswered requests
   343  // that have been sent on the connection.
   344  func (cc *ClientConn) Pending() int {
   345  	cc.lk.Lock()
   346  	defer cc.lk.Unlock()
   347  	return cc.nwritten - cc.nread
   348  }
   349  
   350  // Read reads the next response from the wire. A valid response might be
   351  // returned together with an ErrPersistEOF, which means that the remote
   352  // requested that this be the last request serviced. Read can be called
   353  // concurrently with Write, but not with another Read.
   354  func (cc *ClientConn) Read(req *http.Request) (resp *http.Response, err error) {
   355  	// Retrieve the pipeline ID of this request/response pair
   356  	cc.lk.Lock()
   357  	id, ok := cc.pipereq[req]
   358  	delete(cc.pipereq, req)
   359  	if !ok {
   360  		cc.lk.Unlock()
   361  		return nil, ErrPipeline
   362  	}
   363  	cc.lk.Unlock()
   364  
   365  	// Ensure pipeline order
   366  	cc.pipe.StartResponse(id)
   367  	defer cc.pipe.EndResponse(id)
   368  
   369  	cc.lk.Lock()
   370  	if cc.re != nil {
   371  		defer cc.lk.Unlock()
   372  		return nil, cc.re
   373  	}
   374  	if cc.r == nil { // connection closed by user in the meantime
   375  		defer cc.lk.Unlock()
   376  		return nil, errClosed
   377  	}
   378  	r := cc.r
   379  	lastbody := cc.lastbody
   380  	cc.lastbody = nil
   381  	cc.lk.Unlock()
   382  
   383  	// Make sure body is fully consumed, even if user does not call body.Close
   384  	if lastbody != nil {
   385  		// body.Close is assumed to be idempotent and multiple calls to
   386  		// it should return the error that its first invocation
   387  		// returned.
   388  		err = lastbody.Close()
   389  		if err != nil {
   390  			cc.lk.Lock()
   391  			defer cc.lk.Unlock()
   392  			cc.re = err
   393  			return nil, err
   394  		}
   395  	}
   396  
   397  	resp, err = http.ReadResponse(r, req)
   398  	cc.lk.Lock()
   399  	defer cc.lk.Unlock()
   400  	if err != nil {
   401  		cc.re = err
   402  		return resp, err
   403  	}
   404  	cc.lastbody = resp.Body
   405  
   406  	cc.nread++
   407  
   408  	if resp.Close {
   409  		cc.re = ErrPersistEOF // don't send any more requests
   410  		return resp, cc.re
   411  	}
   412  	return resp, err
   413  }
   414  
   415  // Do is convenience method that writes a request and reads a response.
   416  func (cc *ClientConn) Do(req *http.Request) (resp *http.Response, err error) {
   417  	err = cc.Write(req)
   418  	if err != nil {
   419  		return
   420  	}
   421  	return cc.Read(req)
   422  }