Golang log.Debugf函数代码示例

// Fill calls fn with the available capacity remaining (capacity-fill) and
// fills the bucket with the number of tokens returned by fn. If the remaining
// capacity is 0, Fill returns 0, nil. If the remaining capacity is < 0, Fill
// returns 0, ErrBucketOverflow.
//
// If fn returns an error, it will be returned by Fill along with the remaining
// capacity.
//
// fn is provided the remaining capacity as a soft maximum, fn is allowed to
// use more than the remaining capacity without incurring spillage, though this
// will cause subsequent calls to Fill to return ErrBucketOverflow until the
// next drain.
//
// If the bucket is closed when Fill is called, fn will not be executed and
// Fill will return with an error.
func (b *Bucket) Fill(fn func(int64) (int64, error)) (int64, error) {
	if b.closed() {
		log.Errorf("trafficshape: fill on closed bucket")
		return 0, errFillClosedBucket
	}

	fill := atomic.LoadInt64(&b.fill)
	capacity := atomic.LoadInt64(&b.capacity)

	switch {
	case fill < capacity:
		log.Debugf("trafficshape: under capacity (%d/%d)", fill, capacity)

		n, err := fn(capacity - fill)
		fill = atomic.AddInt64(&b.fill, n)

		return n, err
	case fill > capacity:
		log.Debugf("trafficshape: bucket overflow (%d/%d)", fill, capacity)

		return 0, ErrBucketOverflow
	}

	log.Debugf("trafficshape: bucket full (%d/%d)", fill, capacity)
	return 0, nil
}

func (p *Proxy) connect(req *http.Request) (*http.Response, net.Conn, error) {
	if p.proxyURL != nil {
		log.Debugf("martian: CONNECT with downstream proxy: %s", p.proxyURL.Host)

		conn, err := net.Dial("tcp", p.proxyURL.Host)
		if err != nil {
			return nil, nil, err
		}
		pbw := bufio.NewWriter(conn)
		pbr := bufio.NewReader(conn)

		req.Write(pbw)
		pbw.Flush()

		res, err := http.ReadResponse(pbr, req)
		if err != nil {
			return nil, nil, err
		}

		return res, conn, nil
	}

	log.Debugf("martian: CONNECT to host directly: %s", req.URL.Host)

	conn, err := net.Dial("tcp", req.URL.Host)
	if err != nil {
		return nil, nil, err
	}

	return proxyutil.NewResponse(200, nil, req), conn, nil
}

// loop drains the fill at interval and returns when the bucket is closed.
func (b *Bucket) loop() {
	log.Debugf("trafficshape: started drain loop")
	defer log.Debugf("trafficshape: stopped drain loop")

	for {
		select {
		case t := <-b.t.C:
			atomic.StoreInt64(&b.fill, 0)
			log.Debugf("trafficshape: fill reset @ %s", t)
		case <-b.closec:
			log.Debugf("trafficshape: bucket closed")
			return
		}
	}
}

func (p *Proxy) handleLoop(conn net.Conn) {
	p.conns.Add(1)
	defer p.conns.Done()
	defer conn.Close()

	s, err := newSession()
	if err != nil {
		log.Errorf("martian: failed to create session: %v", err)
		return
	}

	ctx, err := withSession(s)
	if err != nil {
		log.Errorf("martian: failed to create context: %v", err)
		return
	}

	brw := bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))

	for {
		deadline := time.Now().Add(p.timeout)
		conn.SetDeadline(deadline)

		if err := p.handle(ctx, conn, brw); isCloseable(err) {
			log.Debugf("martian: closing connection: %v", conn.RemoteAddr())
			return
		}
	}
}

func (p *Proxy) roundTrip(ctx *Context, req *http.Request) (*http.Response, error) {
	if ctx.SkippingRoundTrip() {
		log.Debugf("martian: skipping round trip")
		return proxyutil.NewResponse(200, nil, req), nil
	}

	return p.roundTripper.RoundTrip(req)
}

// Close stops the drain loop and marks the bucket as closed.
func (b *Bucket) Close() error {
	log.Debugf("trafficshape: closing bucket")

	b.t.Stop()
	close(b.closec)

	return nil
}

// Close closes the read and write buckets along with the underlying listener.
func (l *Listener) Close() error {
	defer log.Debugf("trafficshape: closed read/write buckets and connection")

	l.rb.Close()
	l.wb.Close()

	return l.Listener.Close()
}

// ServeHTTP configures latency and bandwidth constraints.
//
// The "latency" query string parameter accepts a duration string in any format
// supported by time.ParseDuration.
// The "up" and "down" query string parameters accept integers as bits per
// second to be used for read and write throughput.
func (h *Handler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
	log.Debugf("trafficshape: configuration request")

	latency := req.FormValue("latency")
	if latency != "" {
		d, err := time.ParseDuration(latency)
		if err != nil {
			log.Errorf("trafficshape: invalid latency parameter: %v", err)
			http.Error(rw, fmt.Sprintf("invalid duration: %s", latency), 400)
			return
		}

		h.l.SetLatency(d)
	}

	up := req.FormValue("up")
	if up != "" {
		br, err := strconv.ParseInt(up, 10, 64)
		if err != nil {
			log.Errorf("trafficshape: invalid up parameter: %v", err)
			http.Error(rw, fmt.Sprintf("invalid upstream: %s", up), 400)
			return
		}

		h.l.SetWriteBitrate(br)
	}

	down := req.FormValue("down")
	if down != "" {
		br, err := strconv.ParseInt(down, 10, 64)
		if err != nil {
			log.Errorf("trafficshape: invalid down parameter: %v", err)
			http.Error(rw, fmt.Sprintf("invalid downstream: %s", down), 400)
			return
		}

		h.l.SetReadBitrate(br)
	}

	log.Debugf("trafficshape: configured successfully")
}

func (p *Proxy) roundTrip(ctx *session.Context, req *http.Request) (*http.Response, error) {
	if ctx.SkippingRoundTrip() {
		log.Debugf("martian: skipping round trip")
		return proxyutil.NewResponse(200, nil, req), nil
	}

	if tr, ok := p.roundTripper.(*http.Transport); ok {
		tr.Proxy = http.ProxyURL(p.proxyURL)
	}

	return p.roundTripper.RoundTrip(req)
}

// Serve accepts connections from the listener and handles the requests.
func (p *Proxy) Serve(l net.Listener) error {
	defer l.Close()

	var delay time.Duration
	for {
		if p.Closing() {
			return nil
		}

		conn, err := l.Accept()
		if err != nil {
			if nerr, ok := err.(net.Error); ok && nerr.Temporary() {
				if delay == 0 {
					delay = 5 * time.Millisecond
				} else {
					delay *= 2
				}
				if max := time.Second; delay > max {
					delay = max
				}

				log.Debugf("martian: temporary error on accept: %v", err)
				time.Sleep(delay)
				continue
			}

			log.Errorf("martian: failed to accept: %v", err)
			return err
		}
		delay = 0
		log.Debugf("martian: accepted connection from %s", conn.RemoteAddr())

		if tconn, ok := conn.(*net.TCPConn); ok {
			tconn.SetKeepAlive(true)
			tconn.SetKeepAlivePeriod(3 * time.Minute)
		}

		go p.handleLoop(conn)
	}
}

// FillThrottle calls fn with the available capacity remaining (capacity-fill)
// and fills the bucket with the number of tokens returned by fn. If the
// remaining capacity is <= 0, FillThrottle will wait for the next drain before
// running fn.
//
// If fn returns an error, it will be returned by FillThrottle along with the
// number of tokens processed by fn.
//
// fn is provided the remaining capacity as a soft maximum, fn is allowed to
// use more than the remaining capacity without incurring spillage.
//
// If the bucket is closed when FillThrottle is called, or while waiting for
// the next drain, fn will not be executed and FillThrottle will return with an
// error.
func (b *Bucket) FillThrottle(fn func(int64) (int64, error)) (int64, error) {
	for {
		if b.closed() {
			log.Errorf("trafficshape: fill on closed bucket")
			return 0, errFillClosedBucket
		}

		fill := atomic.LoadInt64(&b.fill)
		capacity := atomic.LoadInt64(&b.capacity)

		if fill < capacity {
			log.Debugf("trafficshape: under capacity (%d/%d)", fill, capacity)

			n, err := fn(capacity - fill)
			fill = atomic.AddInt64(&b.fill, n)

			return n, err
		}

		log.Debugf("trafficshape: bucket full (%d/%d)", fill, capacity)
	}
}

// ReadFrom reads data from r until EOF or error, optionally simulating
// connection latency and throttling read throughput based on desired bandwidth
// constraints.
func (c *conn) ReadFrom(r io.Reader) (int64, error) {
	c.ronce.Do(c.sleepLatency)

	var total int64
	for {
		n, err := c.rb.FillThrottle(func(remaining int64) (int64, error) {
			return io.CopyN(c.Conn, r, remaining)
		})

		total += n

		if err == io.EOF {
			log.Debugf("trafficshape: exhausted reader successfully")
			return total, nil
		} else if err != nil {
			log.Errorf("trafficshape: failed copying from reader: %v", err)
			return total, err
		}
	}
}

// Accept waits for and returns the next connection to the listener.
func (l *Listener) Accept() (net.Conn, error) {
	oc, err := l.Listener.Accept()
	if err != nil {
		log.Errorf("trafficshape: failed accepting connection: %v", err)
		return nil, err
	}

	if tconn, ok := oc.(*net.TCPConn); ok {
		log.Debugf("trafficshape: setting keep-alive for TCP connection")
		tconn.SetKeepAlive(true)
		tconn.SetKeepAlivePeriod(3 * time.Minute)
	}

	lc := &conn{
		Conn:    oc,
		latency: l.Latency(),
		rb:      l.rb,
		wb:      l.wb,
	}

	return lc, nil
}

func (c *Config) cert(hostname string) (*tls.Certificate, error) {
	// Remove the port if it exists.
	host, _, err := net.SplitHostPort(hostname)
	if err == nil {
		hostname = host
	}

	c.certmu.RLock()
	tlsc, ok := c.certs[hostname]
	c.certmu.RUnlock()

	if ok {
		log.Debugf("mitm: cache hit for %s", hostname)

		// Check validity of the certificate for hostname match, expiry, etc. In
		// particular, if the cached certificate has expired, create a new one.
		if _, err := tlsc.Leaf.Verify(x509.VerifyOptions{
			DNSName: hostname,
			Roots:   c.roots,
		}); err == nil {
			return tlsc, nil
		}

		log.Debugf("mitm: invalid certificate in cache for %s", hostname)
	}

	log.Debugf("mitm: cache miss for %s", hostname)

	serial, err := rand.Int(rand.Reader, MaxSerialNumber)
	if err != nil {
		return nil, err
	}

	tmpl := &x509.Certificate{
		SerialNumber: serial,
		Subject: pkix.Name{
			CommonName:   hostname,
			Organization: []string{c.org},
		},
		SubjectKeyId:          c.keyID,
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
		NotBefore:             time.Now().Add(-c.validity),
		NotAfter:              time.Now().Add(c.validity),
	}

	if ip := net.ParseIP(hostname); ip != nil {
		tmpl.IPAddresses = []net.IP{ip}
	} else {
		tmpl.DNSNames = []string{hostname}
	}

	raw, err := x509.CreateCertificate(rand.Reader, tmpl, c.ca, c.priv.Public(), c.capriv)
	if err != nil {
		return nil, err
	}

	// Parse certificate bytes so that we have a leaf certificate.
	x509c, err := x509.ParseCertificate(raw)
	if err != nil {
		return nil, err
	}

	tlsc = &tls.Certificate{
		Certificate: [][]byte{raw, c.ca.Raw},
		PrivateKey:  c.priv,
		Leaf:        x509c,
	}

	c.certmu.Lock()
	c.certs[hostname] = tlsc
	c.certmu.Unlock()

	return tlsc, nil
}

// SetCapacity sets the capacity for the bucket and resets the fill to zero.
func (b *Bucket) SetCapacity(capacity int64) {
	log.Debugf("trafficshape: set capacity: %d", capacity)

	atomic.StoreInt64(&b.capacity, capacity)
	atomic.StoreInt64(&b.fill, 0)
}