net.go 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665
  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. package psiphon
  20. import (
  21. "crypto/tls"
  22. "crypto/x509"
  23. "errors"
  24. "fmt"
  25. "io"
  26. "io/ioutil"
  27. "net"
  28. "net/http"
  29. "net/url"
  30. "os"
  31. "reflect"
  32. "sync"
  33. "time"
  34. "github.com/Psiphon-Inc/dns"
  35. "github.com/Psiphon-Inc/ratelimit"
  36. )
  37. const DNS_PORT = 53
  38. // DialConfig contains parameters to determine the behavior
  39. // of a Psiphon dialer (TCPDial, MeekDial, etc.)
  40. type DialConfig struct {
  41. // UpstreamProxyUrl specifies a proxy to connect through.
  42. // E.g., "http://proxyhost:8080"
  43. // "socks5://user:password@proxyhost:1080"
  44. // "socks4a://proxyhost:1080"
  45. // "http://NTDOMAIN\NTUser:password@proxyhost:3375"
  46. //
  47. // Certain tunnel protocols require HTTP CONNECT support
  48. // when a HTTP proxy is specified. If CONNECT is not
  49. // supported, those protocols will not connect.
  50. UpstreamProxyUrl string
  51. ConnectTimeout time.Duration
  52. // PendingConns is used to track and interrupt dials in progress.
  53. // Dials may be interrupted using PendingConns.CloseAll(). Once instantiated,
  54. // a conn is added to pendingConns before the network connect begins and
  55. // removed from pendingConns once the connect succeeds or fails.
  56. // May be nil.
  57. PendingConns *Conns
  58. // BindToDevice parameters are used to exclude connections and
  59. // associated DNS requests from VPN routing.
  60. // When DeviceBinder is set, any underlying socket is
  61. // submitted to the device binding servicebefore connecting.
  62. // The service should bind the socket to a device so that it doesn't route
  63. // through a VPN interface. This service is also used to bind UDP sockets used
  64. // for DNS requests, in which case DnsServerGetter is used to get the
  65. // current active untunneled network DNS server.
  66. DeviceBinder DeviceBinder
  67. DnsServerGetter DnsServerGetter
  68. // UseIndistinguishableTLS specifies whether to try to use an
  69. // alternative stack for TLS. From a circumvention perspective,
  70. // Go's TLS has a distinct fingerprint that may be used for blocking.
  71. // Only applies to TLS connections.
  72. UseIndistinguishableTLS bool
  73. // TrustedCACertificatesFilename specifies a file containing trusted
  74. // CA certs. The file contents should be compatible with OpenSSL's
  75. // SSL_CTX_load_verify_locations.
  76. // Only applies to UseIndistinguishableTLS connections.
  77. TrustedCACertificatesFilename string
  78. // DeviceRegion is the reported region the host device is running in.
  79. // When set, this value may be used, pre-connection, to select performance
  80. // or circumvention optimization strategies for the given region.
  81. DeviceRegion string
  82. // ResolvedIPCallback, when set, is called with the IP address that was
  83. // dialed. This is either the specified IP address in the dial address,
  84. // or the resolved IP address in the case where the dial address is a
  85. // domain name.
  86. // The callback may be invoked by a concurrent goroutine.
  87. ResolvedIPCallback func(string)
  88. }
  89. // NetworkConnectivityChecker defines the interface to the external
  90. // HasNetworkConnectivity provider
  91. type NetworkConnectivityChecker interface {
  92. // TODO: change to bool return value once gobind supports that type
  93. HasNetworkConnectivity() int
  94. }
  95. // DeviceBinder defines the interface to the external BindToDevice provider
  96. type DeviceBinder interface {
  97. BindToDevice(fileDescriptor int) error
  98. }
  99. // DnsServerGetter defines the interface to the external GetDnsServer provider
  100. type DnsServerGetter interface {
  101. GetPrimaryDnsServer() string
  102. GetSecondaryDnsServer() string
  103. }
  104. // HostNameTransformer defines the interface for pluggable hostname
  105. // transformation circumvention strategies.
  106. type HostNameTransformer interface {
  107. TransformHostName(hostname string) (string, bool)
  108. }
  109. // IdentityHostNameTransformer is the default HostNameTransformer, which
  110. // returns the hostname unchanged.
  111. type IdentityHostNameTransformer struct{}
  112. func (IdentityHostNameTransformer) TransformHostName(hostname string) (string, bool) {
  113. return hostname, false
  114. }
  115. // TimeoutError implements the error interface
  116. type TimeoutError struct{}
  117. func (TimeoutError) Error() string { return "timed out" }
  118. func (TimeoutError) Timeout() bool { return true }
  119. func (TimeoutError) Temporary() bool { return true }
  120. // Dialer is a custom dialer compatible with http.Transport.Dial.
  121. type Dialer func(string, string) (net.Conn, error)
  122. // Conns is a synchronized list of Conns that is used to coordinate
  123. // interrupting a set of goroutines establishing connections, or
  124. // close a set of open connections, etc.
  125. // Once the list is closed, no more items may be added to the
  126. // list (unless it is reset).
  127. type Conns struct {
  128. mutex sync.Mutex
  129. isClosed bool
  130. conns map[net.Conn]bool
  131. }
  132. func (conns *Conns) Reset() {
  133. conns.mutex.Lock()
  134. defer conns.mutex.Unlock()
  135. conns.isClosed = false
  136. conns.conns = make(map[net.Conn]bool)
  137. }
  138. func (conns *Conns) Add(conn net.Conn) bool {
  139. conns.mutex.Lock()
  140. defer conns.mutex.Unlock()
  141. if conns.isClosed {
  142. return false
  143. }
  144. if conns.conns == nil {
  145. conns.conns = make(map[net.Conn]bool)
  146. }
  147. conns.conns[conn] = true
  148. return true
  149. }
  150. func (conns *Conns) Remove(conn net.Conn) {
  151. conns.mutex.Lock()
  152. defer conns.mutex.Unlock()
  153. delete(conns.conns, conn)
  154. }
  155. func (conns *Conns) CloseAll() {
  156. conns.mutex.Lock()
  157. defer conns.mutex.Unlock()
  158. conns.isClosed = true
  159. for conn, _ := range conns.conns {
  160. conn.Close()
  161. }
  162. conns.conns = make(map[net.Conn]bool)
  163. }
  164. // LocalProxyRelay sends to remoteConn bytes received from localConn,
  165. // and sends to localConn bytes received from remoteConn.
  166. func LocalProxyRelay(proxyType string, localConn, remoteConn net.Conn) {
  167. copyWaitGroup := new(sync.WaitGroup)
  168. copyWaitGroup.Add(1)
  169. go func() {
  170. defer copyWaitGroup.Done()
  171. _, err := io.Copy(localConn, remoteConn)
  172. if err != nil {
  173. err = fmt.Errorf("Relay failed: %s", ContextError(err))
  174. NoticeLocalProxyError(proxyType, err)
  175. }
  176. }()
  177. _, err := io.Copy(remoteConn, localConn)
  178. if err != nil {
  179. err = fmt.Errorf("Relay failed: %s", ContextError(err))
  180. NoticeLocalProxyError(proxyType, err)
  181. }
  182. copyWaitGroup.Wait()
  183. }
  184. // WaitForNetworkConnectivity uses a NetworkConnectivityChecker to
  185. // periodically check for network connectivity. It returns true if
  186. // no NetworkConnectivityChecker is provided (waiting is disabled)
  187. // or when NetworkConnectivityChecker.HasNetworkConnectivity()
  188. // indicates connectivity. It waits and polls the checker once a second.
  189. // If any stop is broadcast, false is returned immediately.
  190. func WaitForNetworkConnectivity(
  191. connectivityChecker NetworkConnectivityChecker, stopBroadcasts ...<-chan struct{}) bool {
  192. if connectivityChecker == nil || 1 == connectivityChecker.HasNetworkConnectivity() {
  193. return true
  194. }
  195. NoticeInfo("waiting for network connectivity")
  196. ticker := time.NewTicker(1 * time.Second)
  197. for {
  198. if 1 == connectivityChecker.HasNetworkConnectivity() {
  199. return true
  200. }
  201. selectCases := make([]reflect.SelectCase, 1+len(stopBroadcasts))
  202. selectCases[0] = reflect.SelectCase{
  203. Dir: reflect.SelectRecv, Chan: reflect.ValueOf(ticker.C)}
  204. for i, stopBroadcast := range stopBroadcasts {
  205. selectCases[i+1] = reflect.SelectCase{
  206. Dir: reflect.SelectRecv, Chan: reflect.ValueOf(stopBroadcast)}
  207. }
  208. chosen, _, ok := reflect.Select(selectCases)
  209. if chosen == 0 && ok {
  210. // Ticker case, so check again
  211. } else {
  212. // Stop case
  213. return false
  214. }
  215. }
  216. }
  217. // ResolveIP uses a custom dns stack to make a DNS query over the
  218. // given TCP or UDP conn. This is used, e.g., when we need to ensure
  219. // that a DNS connection bypasses a VPN interface (BindToDevice) or
  220. // when we need to ensure that a DNS connection is tunneled.
  221. // Caller must set timeouts or interruptibility as required for conn.
  222. func ResolveIP(host string, conn net.Conn) (addrs []net.IP, ttls []time.Duration, err error) {
  223. // Send the DNS query
  224. dnsConn := &dns.Conn{Conn: conn}
  225. defer dnsConn.Close()
  226. query := new(dns.Msg)
  227. query.SetQuestion(dns.Fqdn(host), dns.TypeA)
  228. query.RecursionDesired = true
  229. dnsConn.WriteMsg(query)
  230. // Process the response
  231. response, err := dnsConn.ReadMsg()
  232. if err != nil {
  233. return nil, nil, ContextError(err)
  234. }
  235. addrs = make([]net.IP, 0)
  236. ttls = make([]time.Duration, 0)
  237. for _, answer := range response.Answer {
  238. if a, ok := answer.(*dns.A); ok {
  239. addrs = append(addrs, a.A)
  240. ttl := time.Duration(a.Hdr.Ttl) * time.Second
  241. ttls = append(ttls, ttl)
  242. }
  243. }
  244. return addrs, ttls, nil
  245. }
  246. // MakeUntunneledHttpsClient returns a net/http.Client which is
  247. // configured to use custom dialing features -- including BindToDevice,
  248. // UseIndistinguishableTLS, etc. -- for a specific HTTPS request URL.
  249. // If verifyLegacyCertificate is not nil, it's used for certificate
  250. // verification.
  251. // Because UseIndistinguishableTLS requires a hack to work with
  252. // net/http, MakeUntunneledHttpClient may return a modified request URL
  253. // to be used. Callers should always use this return value to make
  254. // requests, not the input value.
  255. func MakeUntunneledHttpsClient(
  256. dialConfig *DialConfig,
  257. verifyLegacyCertificate *x509.Certificate,
  258. requestUrl string,
  259. requestTimeout time.Duration) (*http.Client, string, error) {
  260. // Change the scheme to "http"; otherwise http.Transport will try to do
  261. // another TLS handshake inside the explicit TLS session. Also need to
  262. // force an explicit port, as the default for "http", 80, won't talk TLS.
  263. urlComponents, err := url.Parse(requestUrl)
  264. if err != nil {
  265. return nil, "", ContextError(err)
  266. }
  267. urlComponents.Scheme = "http"
  268. host, port, err := net.SplitHostPort(urlComponents.Host)
  269. if err != nil {
  270. // Assume there's no port
  271. host = urlComponents.Host
  272. port = ""
  273. }
  274. if port == "" {
  275. port = "443"
  276. }
  277. urlComponents.Host = net.JoinHostPort(host, port)
  278. // Note: IndistinguishableTLS mode doesn't support VerifyLegacyCertificate
  279. useIndistinguishableTLS := dialConfig.UseIndistinguishableTLS && verifyLegacyCertificate == nil
  280. dialer := NewCustomTLSDialer(
  281. // Note: when verifyLegacyCertificate is not nil, some
  282. // of the other CustomTLSConfig is overridden.
  283. &CustomTLSConfig{
  284. Dial: NewTCPDialer(dialConfig),
  285. VerifyLegacyCertificate: verifyLegacyCertificate,
  286. SNIServerName: host,
  287. SkipVerify: false,
  288. UseIndistinguishableTLS: useIndistinguishableTLS,
  289. TrustedCACertificatesFilename: dialConfig.TrustedCACertificatesFilename,
  290. })
  291. transport := &http.Transport{
  292. Dial: dialer,
  293. }
  294. httpClient := &http.Client{
  295. Timeout: requestTimeout,
  296. Transport: transport,
  297. }
  298. return httpClient, urlComponents.String(), nil
  299. }
  300. // MakeTunneledHttpClient returns a net/http.Client which is
  301. // configured to use custom dialing features including tunneled
  302. // dialing and, optionally, UseTrustedCACertificatesForStockTLS.
  303. // Unlike MakeUntunneledHttpsClient and makePsiphonHttpsClient,
  304. // This http.Client uses stock TLS and no scheme transformation
  305. // hack is required.
  306. func MakeTunneledHttpClient(
  307. config *Config,
  308. tunnel *Tunnel,
  309. requestTimeout time.Duration) (*http.Client, error) {
  310. tunneledDialer := func(_, addr string) (conn net.Conn, err error) {
  311. return tunnel.sshClient.Dial("tcp", addr)
  312. }
  313. transport := &http.Transport{
  314. Dial: tunneledDialer,
  315. ResponseHeaderTimeout: requestTimeout,
  316. }
  317. if config.UseTrustedCACertificatesForStockTLS {
  318. if config.TrustedCACertificatesFilename == "" {
  319. return nil, ContextError(errors.New(
  320. "UseTrustedCACertificatesForStockTLS requires TrustedCACertificatesFilename"))
  321. }
  322. rootCAs := x509.NewCertPool()
  323. certData, err := ioutil.ReadFile(config.TrustedCACertificatesFilename)
  324. if err != nil {
  325. return nil, ContextError(err)
  326. }
  327. rootCAs.AppendCertsFromPEM(certData)
  328. transport.TLSClientConfig = &tls.Config{RootCAs: rootCAs}
  329. }
  330. return &http.Client{
  331. Transport: transport,
  332. Timeout: requestTimeout,
  333. }, nil
  334. }
  335. // MakeDownloadHttpClient is a resusable helper that sets up a
  336. // http.Client for use either untunneled or through a tunnel.
  337. // See MakeUntunneledHttpsClient for a note about request URL
  338. // rewritting.
  339. func MakeDownloadHttpClient(
  340. config *Config,
  341. tunnel *Tunnel,
  342. untunneledDialConfig *DialConfig,
  343. requestUrl string,
  344. requestTimeout time.Duration) (*http.Client, string, error) {
  345. var httpClient *http.Client
  346. var err error
  347. if tunnel != nil {
  348. httpClient, err = MakeTunneledHttpClient(config, tunnel, requestTimeout)
  349. if err != nil {
  350. return nil, "", ContextError(err)
  351. }
  352. } else {
  353. httpClient, requestUrl, err = MakeUntunneledHttpsClient(
  354. untunneledDialConfig, nil, requestUrl, requestTimeout)
  355. if err != nil {
  356. return nil, "", ContextError(err)
  357. }
  358. }
  359. return httpClient, requestUrl, nil
  360. }
  361. // ResumeDownload is a resuable helper that downloads requestUrl via the
  362. // httpClient, storing the result in downloadFilename when the download is
  363. // complete. Intermediate, partial downloads state is stored in
  364. // downloadFilename.part and downloadFilename.part.etag.
  365. // Any existing downloadFilename file will be overwritten.
  366. //
  367. // In the case where the remote object has change while a partial download
  368. // is to be resumed, the partial state is reset and resumeDownload fails.
  369. // The caller must restart the download.
  370. //
  371. // When ifNoneMatchETag is specified, no download is made if the remote
  372. // object has the same ETag. ifNoneMatchETag has an effect only when no
  373. // partial download is in progress.
  374. //
  375. func ResumeDownload(
  376. httpClient *http.Client,
  377. requestUrl string,
  378. downloadFilename string,
  379. ifNoneMatchETag string) (int64, string, error) {
  380. partialFilename := fmt.Sprintf("%s.part", downloadFilename)
  381. partialETagFilename := fmt.Sprintf("%s.part.etag", downloadFilename)
  382. file, err := os.OpenFile(partialFilename, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
  383. if err != nil {
  384. return 0, "", ContextError(err)
  385. }
  386. defer file.Close()
  387. fileInfo, err := file.Stat()
  388. if err != nil {
  389. return 0, "", ContextError(err)
  390. }
  391. // A partial download should have an ETag which is to be sent with the
  392. // Range request to ensure that the source object is the same as the
  393. // one that is partially downloaded.
  394. var partialETag []byte
  395. if fileInfo.Size() > 0 {
  396. partialETag, err = ioutil.ReadFile(partialETagFilename)
  397. // When the ETag can't be loaded, delete the partial download. To keep the
  398. // code simple, there is no immediate, inline retry here, on the assumption
  399. // that the controller's upgradeDownloader will shortly call DownloadUpgrade
  400. // again.
  401. if err != nil {
  402. os.Remove(partialFilename)
  403. os.Remove(partialETagFilename)
  404. return 0, "", ContextError(
  405. fmt.Errorf("failed to load partial download ETag: %s", err))
  406. }
  407. }
  408. request, err := http.NewRequest("GET", requestUrl, nil)
  409. if err != nil {
  410. return 0, "", ContextError(err)
  411. }
  412. request.Header.Add("Range", fmt.Sprintf("bytes=%d-", fileInfo.Size()))
  413. if partialETag != nil {
  414. // Note: not using If-Range, since not all host servers support it.
  415. // Using If-Match means we need to check for status code 412 and reset
  416. // when the ETag has changed since the last partial download.
  417. request.Header.Add("If-Match", string(partialETag))
  418. } else if ifNoneMatchETag != "" {
  419. // Can't specify both If-Match and If-None-Match. Behavior is undefined.
  420. // https://www.w3.org/Protocols/rfc2616/rfc2616-sec14.html#sec14.26
  421. // So for downloaders that store an ETag and wish to use that to prevent
  422. // redundant downloads, that ETag is sent as If-None-Match in the case
  423. // where a partial download is not in progress. When a partial download
  424. // is in progress, the partial ETag is sent as If-Match: either that's
  425. // a version that was never fully received, or it's no longer current in
  426. // which case the response will be StatusPreconditionFailed, the partial
  427. // download will be discarded, and then the next retry will use
  428. // If-None-Match.
  429. // Note: in this case, fileInfo.Size() == 0
  430. request.Header.Add("If-None-Match", ifNoneMatchETag)
  431. }
  432. response, err := httpClient.Do(request)
  433. // The resumeable download may ask for bytes past the resource range
  434. // since it doesn't store the "completed download" state. In this case,
  435. // the HTTP server returns 416. Otherwise, we expect 206. We may also
  436. // receive 412 on ETag mismatch.
  437. if err == nil &&
  438. (response.StatusCode != http.StatusPartialContent &&
  439. response.StatusCode != http.StatusRequestedRangeNotSatisfiable &&
  440. response.StatusCode != http.StatusPreconditionFailed &&
  441. response.StatusCode != http.StatusNotModified) {
  442. response.Body.Close()
  443. err = fmt.Errorf("unexpected response status code: %d", response.StatusCode)
  444. }
  445. if err != nil {
  446. return 0, "", ContextError(err)
  447. }
  448. defer response.Body.Close()
  449. responseETag := response.Header.Get("ETag")
  450. if response.StatusCode == http.StatusPreconditionFailed {
  451. // When the ETag no longer matches, delete the partial download. As above,
  452. // simply failing and relying on the caller's retry schedule.
  453. os.Remove(partialFilename)
  454. os.Remove(partialETagFilename)
  455. return 0, "", ContextError(errors.New("partial download ETag mismatch"))
  456. } else if response.StatusCode == http.StatusNotModified {
  457. // This status code is possible in the "If-None-Match" case. Don't leave
  458. // any partial download in progress. Caller should check that responseETag
  459. // matches ifNoneMatchETag.
  460. os.Remove(partialFilename)
  461. os.Remove(partialETagFilename)
  462. return 0, responseETag, nil
  463. }
  464. // Not making failure to write ETag file fatal, in case the entire download
  465. // succeeds in this one request.
  466. ioutil.WriteFile(partialETagFilename, []byte(responseETag), 0600)
  467. // A partial download occurs when this copy is interrupted. The io.Copy
  468. // will fail, leaving a partial download in place (.part and .part.etag).
  469. n, err := io.Copy(NewSyncFileWriter(file), response.Body)
  470. // From this point, n bytes are indicated as downloaded, even if there is
  471. // an error; the caller may use this to report partial download progress.
  472. if err != nil {
  473. return n, "", ContextError(err)
  474. }
  475. // Ensure the file is flushed to disk. The deferred close
  476. // will be a noop when this succeeds.
  477. err = file.Close()
  478. if err != nil {
  479. return n, "", ContextError(err)
  480. }
  481. // Remove if exists, to enable rename
  482. os.Remove(downloadFilename)
  483. err = os.Rename(partialFilename, downloadFilename)
  484. if err != nil {
  485. return n, "", ContextError(err)
  486. }
  487. os.Remove(partialETagFilename)
  488. return n, responseETag, nil
  489. }
  490. // IPAddressFromAddr is a helper which extracts an IP address
  491. // from a net.Addr or returns "" if there is no IP address.
  492. func IPAddressFromAddr(addr net.Addr) string {
  493. ipAddress := ""
  494. if addr != nil {
  495. host, _, err := net.SplitHostPort(addr.String())
  496. if err == nil {
  497. ipAddress = host
  498. }
  499. }
  500. return ipAddress
  501. }
  502. // IdleTimeoutConn wraps a net.Conn and sets an initial ReadDeadline. The
  503. // deadline is extended whenever data is received from the connection.
  504. // Optionally, IdleTimeoutConn will also extend the deadline when data is
  505. // written to the connection.
  506. type IdleTimeoutConn struct {
  507. net.Conn
  508. deadline time.Duration
  509. resetOnWrite bool
  510. }
  511. func NewIdleTimeoutConn(
  512. conn net.Conn, deadline time.Duration, resetOnWrite bool) *IdleTimeoutConn {
  513. conn.SetReadDeadline(time.Now().Add(deadline))
  514. return &IdleTimeoutConn{
  515. Conn: conn,
  516. deadline: deadline,
  517. resetOnWrite: resetOnWrite,
  518. }
  519. }
  520. func (conn *IdleTimeoutConn) Read(buffer []byte) (int, error) {
  521. n, err := conn.Conn.Read(buffer)
  522. if err == nil {
  523. conn.Conn.SetReadDeadline(time.Now().Add(conn.deadline))
  524. }
  525. return n, err
  526. }
  527. func (conn *IdleTimeoutConn) Write(buffer []byte) (int, error) {
  528. n, err := conn.Conn.Write(buffer)
  529. if err == nil && conn.resetOnWrite {
  530. conn.Conn.SetReadDeadline(time.Now().Add(conn.deadline))
  531. }
  532. return n, err
  533. }
  534. // ThrottledConn wraps a net.Conn with read and write rate limiters.
  535. // Rates are specified as bytes per second. The underlying rate limiter
  536. // uses the token bucket algorithm to calculate delay times for read
  537. // and write operations. Specify limit values of 0 set no limit.
  538. type ThrottledConn struct {
  539. net.Conn
  540. reader io.Reader
  541. writer io.Writer
  542. }
  543. func NewThrottledConn(
  544. conn net.Conn,
  545. limitReadBytesPerSecond, limitWriteBytesPerSecond int64) *ThrottledConn {
  546. var reader io.Reader
  547. if limitReadBytesPerSecond == 0 {
  548. reader = conn
  549. } else {
  550. reader = ratelimit.Reader(conn,
  551. ratelimit.NewBucketWithRate(
  552. float64(limitReadBytesPerSecond), limitReadBytesPerSecond))
  553. }
  554. var writer io.Writer
  555. if limitWriteBytesPerSecond == 0 {
  556. writer = conn
  557. } else {
  558. writer = ratelimit.Writer(conn,
  559. ratelimit.NewBucketWithRate(
  560. float64(limitWriteBytesPerSecond), limitWriteBytesPerSecond))
  561. }
  562. return &ThrottledConn{
  563. Conn: conn,
  564. reader: reader,
  565. writer: writer,
  566. }
  567. }
  568. func (conn *ThrottledConn) Read(buffer []byte) (int, error) {
  569. return conn.reader.Read(buffer)
  570. }
  571. func (conn *ThrottledConn) Write(buffer []byte) (int, error) {
  572. return conn.writer.Write(buffer)
  573. }