github.com/psiphon-Labs/psiphon-tunnel-core@v2.0.28+incompatible/psiphon/common/throttled_test.go (about)

     1  /*
     2   * Copyright (c) 2016, 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 common
    21  
    22  import (
    23  	"bytes"
    24  	"fmt"
    25  	"io/ioutil"
    26  	"math"
    27  	"net"
    28  	"net/http"
    29  	"testing"
    30  	"time"
    31  )
    32  
    33  const (
    34  	serverAddress = "127.0.0.1:8081"
    35  	testDataSize  = 10 * 1024 * 1024 // 10 MB
    36  )
    37  
    38  func TestThrottledConnRates(t *testing.T) {
    39  
    40  	runRateLimitsTest(t, RateLimits{
    41  		ReadUnthrottledBytes:  0,
    42  		ReadBytesPerSecond:    0,
    43  		WriteUnthrottledBytes: 0,
    44  		WriteBytesPerSecond:   0,
    45  	})
    46  
    47  	runRateLimitsTest(t, RateLimits{
    48  		ReadUnthrottledBytes:  0,
    49  		ReadBytesPerSecond:    5 * 1024 * 1024,
    50  		WriteUnthrottledBytes: 0,
    51  		WriteBytesPerSecond:   5 * 1024 * 1024,
    52  	})
    53  
    54  	runRateLimitsTest(t, RateLimits{
    55  		ReadUnthrottledBytes:  0,
    56  		ReadBytesPerSecond:    5 * 1024 * 1024,
    57  		WriteUnthrottledBytes: 0,
    58  		WriteBytesPerSecond:   1024 * 1024,
    59  	})
    60  
    61  	runRateLimitsTest(t, RateLimits{
    62  		ReadUnthrottledBytes:  0,
    63  		ReadBytesPerSecond:    2 * 1024 * 1024,
    64  		WriteUnthrottledBytes: 0,
    65  		WriteBytesPerSecond:   2 * 1024 * 1024,
    66  	})
    67  
    68  	runRateLimitsTest(t, RateLimits{
    69  		ReadUnthrottledBytes:  0,
    70  		ReadBytesPerSecond:    1024 * 1024,
    71  		WriteUnthrottledBytes: 0,
    72  		WriteBytesPerSecond:   1024 * 1024,
    73  	})
    74  
    75  	// This test takes > 1 min to run, so disabled for now
    76  	/*
    77  		runRateLimitsTest(t, RateLimits{
    78  			ReadUnthrottledBytes: 0,
    79  			ReadBytesPerSecond: 1024 * 1024 / 8,
    80  			WriteUnthrottledBytes:   0,
    81  			WriteBytesPerSecond:   1024 * 1024 / 8,
    82  		})
    83  	*/
    84  }
    85  
    86  func runRateLimitsTest(t *testing.T, rateLimits RateLimits) {
    87  
    88  	// Run a local HTTP server which serves large chunks of data
    89  
    90  	go func() {
    91  
    92  		handler := func(w http.ResponseWriter, r *http.Request) {
    93  			_, _ = ioutil.ReadAll(r.Body)
    94  			testData, _ := MakeSecureRandomBytes(testDataSize)
    95  			w.Write(testData)
    96  		}
    97  
    98  		server := &http.Server{
    99  			Addr:    serverAddress,
   100  			Handler: http.HandlerFunc(handler),
   101  		}
   102  
   103  		server.ListenAndServe()
   104  	}()
   105  
   106  	// TODO: properly synchronize with server startup
   107  	time.Sleep(1 * time.Second)
   108  
   109  	// Set up a HTTP client with a throttled connection
   110  
   111  	throttledDial := func(network, addr string) (net.Conn, error) {
   112  		conn, err := net.Dial(network, addr)
   113  		if err != nil {
   114  			return conn, err
   115  		}
   116  		return NewThrottledConn(conn, rateLimits), nil
   117  	}
   118  
   119  	client := &http.Client{
   120  		Transport: &http.Transport{
   121  			Dial: throttledDial,
   122  		},
   123  	}
   124  
   125  	// Upload and download a large chunk of data, and time it
   126  
   127  	testData, _ := MakeSecureRandomBytes(testDataSize)
   128  	requestBody := bytes.NewReader(testData)
   129  
   130  	startTime := time.Now()
   131  
   132  	response, err := client.Post("http://"+serverAddress, "application/octet-stream", requestBody)
   133  	if err == nil && response.StatusCode != http.StatusOK {
   134  		response.Body.Close()
   135  		err = fmt.Errorf("unexpected response code: %d", response.StatusCode)
   136  	}
   137  	if err != nil {
   138  		t.Fatalf("request failed: %s", err)
   139  	}
   140  	defer response.Body.Close()
   141  
   142  	// Test: elapsed upload time must reflect rate limit
   143  
   144  	checkElapsedTime(t, testDataSize, rateLimits.WriteBytesPerSecond, time.Since(startTime))
   145  
   146  	startTime = time.Now()
   147  
   148  	body, err := ioutil.ReadAll(response.Body)
   149  	if err != nil {
   150  		t.Fatalf("read response failed: %s", err)
   151  	}
   152  	if len(body) != testDataSize {
   153  		t.Fatalf("unexpected response size: %d", len(body))
   154  	}
   155  
   156  	// Test: elapsed download time must reflect rate limit
   157  
   158  	checkElapsedTime(t, testDataSize, rateLimits.ReadBytesPerSecond, time.Since(startTime))
   159  }
   160  
   161  func checkElapsedTime(t *testing.T, dataSize int, rateLimit int64, duration time.Duration) {
   162  
   163  	// With no rate limit, should finish under a couple seconds
   164  	floorElapsedTime := 0 * time.Second
   165  	ceilingElapsedTime := 2 * time.Second
   166  
   167  	if rateLimit != 0 {
   168  		// With rate limit, should finish within a couple seconds or so of data size / bytes-per-second;
   169  		// won't be exact due to request overhead and approximations in "ratelimit" package
   170  		expectedElapsedTime := float64(testDataSize) / float64(rateLimit)
   171  		floorElapsedTime = time.Duration(int64(math.Floor(expectedElapsedTime))) * time.Second
   172  		floorElapsedTime -= 1500 * time.Millisecond
   173  		if floorElapsedTime < 0 {
   174  			floorElapsedTime = 0
   175  		}
   176  		ceilingElapsedTime = time.Duration(int64(math.Ceil(expectedElapsedTime))) * time.Second
   177  		ceilingElapsedTime += 1500 * time.Millisecond
   178  	}
   179  
   180  	t.Logf(
   181  		"\ndata size: %d\nrate limit: %d\nelapsed time: %s\nexpected time: [%s,%s]\n\n",
   182  		dataSize,
   183  		rateLimit,
   184  		duration,
   185  		floorElapsedTime,
   186  		ceilingElapsedTime)
   187  
   188  	if duration < floorElapsedTime {
   189  		t.Errorf("unexpected duration: %s < %s", duration, floorElapsedTime)
   190  	}
   191  
   192  	if duration > ceilingElapsedTime {
   193  		t.Errorf("unexpected duration: %s > %s", duration, ceilingElapsedTime)
   194  	}
   195  }
   196  
   197  func TestThrottledConnClose(t *testing.T) {
   198  
   199  	rateLimits := RateLimits{
   200  		ReadBytesPerSecond:  1,
   201  		WriteBytesPerSecond: 1,
   202  	}
   203  
   204  	n := 4
   205  	b := make([]byte, n+1)
   206  
   207  	throttledConn := NewThrottledConn(&testConn{}, rateLimits)
   208  
   209  	now := time.Now()
   210  	_, err := throttledConn.Read(b)
   211  	elapsed := time.Since(now)
   212  	if err != nil || elapsed < time.Duration(n)*time.Second {
   213  		t.Errorf("unexpected interrupted read: %s, %v", elapsed, err)
   214  	}
   215  
   216  	now = time.Now()
   217  	go func() {
   218  		time.Sleep(500 * time.Millisecond)
   219  		throttledConn.Close()
   220  	}()
   221  	_, err = throttledConn.Read(b)
   222  	elapsed = time.Since(now)
   223  	if elapsed > 1*time.Second {
   224  		t.Errorf("unexpected uninterrupted read: %s, %v", elapsed, err)
   225  	}
   226  
   227  	throttledConn = NewThrottledConn(&testConn{}, rateLimits)
   228  
   229  	now = time.Now()
   230  	_, err = throttledConn.Write(b)
   231  	elapsed = time.Since(now)
   232  	if err != nil || elapsed < time.Duration(n)*time.Second {
   233  		t.Errorf("unexpected interrupted write: %s, %v", elapsed, err)
   234  	}
   235  
   236  	now = time.Now()
   237  	go func() {
   238  		time.Sleep(500 * time.Millisecond)
   239  		throttledConn.Close()
   240  	}()
   241  	_, err = throttledConn.Write(b)
   242  	elapsed = time.Since(now)
   243  	if elapsed > 1*time.Second {
   244  		t.Errorf("unexpected uninterrupted write: %s, %v", elapsed, err)
   245  	}
   246  }
   247  
   248  type testConn struct {
   249  }
   250  
   251  func (conn *testConn) Read(b []byte) (n int, err error) {
   252  	return len(b), nil
   253  }
   254  
   255  func (conn *testConn) Write(b []byte) (n int, err error) {
   256  	return len(b), nil
   257  }
   258  
   259  func (conn *testConn) Close() error {
   260  	return nil
   261  }
   262  
   263  func (conn *testConn) LocalAddr() net.Addr {
   264  	return nil
   265  }
   266  
   267  func (conn *testConn) RemoteAddr() net.Addr {
   268  	return nil
   269  }
   270  
   271  func (conn *testConn) SetDeadline(t time.Time) error {
   272  	return nil
   273  }
   274  
   275  func (conn *testConn) SetReadDeadline(t time.Time) error {
   276  	return nil
   277  }
   278  
   279  func (conn *testConn) SetWriteDeadline(t time.Time) error {
   280  	return nil
   281  }