github.com/psiphon-labs/psiphon-tunnel-core@v2.0.28+incompatible/psiphon/transferstats/conn.go (about)

     1  /*
     2   * Copyright (c) 2015, Psiphon Inc.
     3   * All rights reserved.
     4   *
     5   * This program is free software: you can redistribute it and/or modify
     6   * it under the terms of the GNU General Public License as published by
     7   * the Free Software Foundation, either version 3 of the License, or
     8   * (at your option) any later version.
     9   *
    10   * This program is distributed in the hope that it will be useful,
    11   * but WITHOUT ANY WARRANTY; without even the implied warranty of
    12   * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    13   * GNU General Public License for more details.
    14   *
    15   * You should have received a copy of the GNU General Public License
    16   * along with this program.  If not, see <http://www.gnu.org/licenses/>.
    17   *
    18   */
    19  
    20  // Package transferstats counts and keeps track of session stats. These are
    21  // per-domain bytes transferred and total bytes transferred.
    22  package transferstats
    23  
    24  /*
    25  Assumption: The same connection will not be used to access different hostnames
    26  	(even if, say, those hostnames map to the same server). If this does occur, we
    27  	will mis-attribute some bytes.
    28  Assumption: Enough of the first HTTP will be present in the first Write() call
    29  	for us to a) recognize that it is HTTP, and b) parse the hostname.
    30  		- If this turns out to not be generally true we will need to add buffering.
    31  */
    32  
    33  import (
    34  	"net"
    35  	"sync/atomic"
    36  )
    37  
    38  // Conn is to be used as an intermediate link in a chain of net.Conn objects.
    39  // It inspects requests and responses and derives stats from them.
    40  type Conn struct {
    41  	net.Conn
    42  	serverID       string
    43  	firstWrite     int32
    44  	hostnameParsed int32
    45  	hostname       string
    46  	regexps        *Regexps
    47  }
    48  
    49  // NewConn creates a Conn. serverID can be anything that uniquely
    50  // identifies the server; it will be passed to TakeOutStatsForServer() when
    51  // retrieving the accumulated stats.
    52  func NewConn(nextConn net.Conn, serverID string, regexps *Regexps) *Conn {
    53  	return &Conn{
    54  		Conn:           nextConn,
    55  		serverID:       serverID,
    56  		firstWrite:     1,
    57  		hostnameParsed: 0,
    58  		regexps:        regexps,
    59  	}
    60  }
    61  
    62  func (conn *Conn) isRecordingHostBytes() bool {
    63  
    64  	// When there are no regexes, no per-host bytes stats will be
    65  	// recorded, including no "(OTHER)" category. In this case, it's
    66  	// expected that there will be no data to send in any status
    67  	// request.
    68  
    69  	return conn.regexps != nil && len(*conn.regexps) > 0
    70  }
    71  
    72  // Write is called when requests are being written out through the tunnel to
    73  // the remote server.
    74  func (conn *Conn) Write(buffer []byte) (n int, err error) {
    75  	// First pass the data down the chain.
    76  	n, err = conn.Conn.Write(buffer)
    77  
    78  	// Count stats before we check the error condition. It could happen that the
    79  	// buffer was partially written and then an error occurred.
    80  	if n > 0 {
    81  
    82  		// If this is the first request, try to determine the hostname to associate
    83  		// with this connection. Skip parsing when not recording per-host bytes, as
    84  		// the hostname isn't used in this case.
    85  
    86  		if conn.isRecordingHostBytes() && atomic.CompareAndSwapInt32(&conn.firstWrite, 1, 0) {
    87  
    88  			hostname, ok := getHostname(buffer)
    89  			if ok {
    90  				// Get the hostname value that will be stored in stats by
    91  				// regexing the real hostname.
    92  				conn.hostname = regexHostname(hostname, conn.regexps)
    93  				atomic.StoreInt32(&conn.hostnameParsed, 1)
    94  			}
    95  		}
    96  
    97  		recordStat(&statsUpdate{
    98  			conn.serverID,
    99  			conn.hostname,
   100  			int64(n),
   101  			0},
   102  			conn.isRecordingHostBytes(),
   103  			false)
   104  	}
   105  
   106  	return
   107  }
   108  
   109  // Read is called when responses to requests are being read from the remote server.
   110  func (conn *Conn) Read(buffer []byte) (n int, err error) {
   111  	n, err = conn.Conn.Read(buffer)
   112  
   113  	var hostname string
   114  	if atomic.LoadInt32(&conn.hostnameParsed) == 1 {
   115  		hostname = conn.hostname
   116  	} else {
   117  		hostname = ""
   118  	}
   119  
   120  	// Count bytes without checking the error condition. It could happen that the
   121  	// buffer was partially read and then an error occurred.
   122  	recordStat(&statsUpdate{
   123  		conn.serverID,
   124  		hostname,
   125  		0,
   126  		int64(n)},
   127  		conn.isRecordingHostBytes(),
   128  		false)
   129  
   130  	return
   131  }