<commit_before>package kafka

import (
	"context"
	"crypto/tls"
	"io"
	"net"
	"strconv"
	"strings"
	"time"

	"github.com/segmentio/kafka-go/sasl"
)

// The Dialer type mirrors the net.Dialer API but is designed to open kafka
// connections instead of raw network connections.
type Dialer struct {
	// Unique identifier for client connections established by this Dialer.
	ClientID string

	// Timeout is the maximum amount of time a dial will wait for a connect to
	// complete. If Deadline is also set, it may fail earlier.
	//
	// The default is no timeout.
	//
	// When dialing a name with multiple IP addresses, the timeout may be
	// divided between them.
	//
	// With or without a timeout, the operating system may impose its own
	// earlier timeout. For instance, TCP timeouts are often around 3 minutes.
	Timeout time.Duration

	// Deadline is the absolute point in time after which dials will fail.
	// If Timeout is set, it may fail earlier.
	// Zero means no deadline, or dependent on the operating system as with the
	// Timeout option.
	Deadline time.Time

	// LocalAddr is the local address to use when dialing an address.
	// The address must be of a compatible type for the network being dialed.
	// If nil, a local address is automatically chosen.
	LocalAddr net.Addr

	// DualStack enables RFC 6555-compliant "Happy Eyeballs" dialing when the
	// network is "tcp" and the destination is a host name with both IPv4 and
	// IPv6 addresses. This allows a client to tolerate networks where one
	// address family is silently broken.
	DualStack bool

	// FallbackDelay specifies the length of time to wait before spawning a
	// fallback connection, when DualStack is enabled.
	// If zero, a default delay of 300ms is used.
	FallbackDelay time.Duration

	// KeepAlive specifies the keep-alive period for an active network
	// connection.
	// If zero, keep-alives are not enabled. Network protocols that do not
	// support keep-alives ignore this field.
	KeepAlive time.Duration

	// Resolver optionally specifies an alternate resolver to use.
	Resolver Resolver

	// TLS enables Dialer to open secure connections.  If nil, standard net.Conn
	// will be used.
	TLS *tls.Config

	// SASLMechanism configures the Dialer to use SASL authentication.  If nil,
	// no authentication will be performed.
	SASLMechanism sasl.Mechanism

	// The transactional id to use for transactional delivery. Idempotent
	// deliver should be enabled if transactional id is configured.
	// For more details look at transactional.id description here: http://kafka.apache.org/documentation.html#producerconfigs
	// Empty string means that the connection will be non-transactional.
	TransactionalID string
}

// Dial connects to the address on the named network.
func (d *Dialer) Dial(network string, address string) (*Conn, error) {
	return d.DialContext(context.Background(), network, address)
}

// DialContext connects to the address on the named network using the provided
// context.
//
// The provided Context must be non-nil. If the context expires before the
// connection is complete, an error is returned. Once successfully connected,
// any expiration of the context will not affect the connection.
//
// When using TCP, and the host in the address parameter resolves to multiple
// network addresses, any dial timeout (from d.Timeout or ctx) is spread over
// each consecutive dial, such that each is given an appropriate fraction of the
// time to connect. For example, if a host has 4 IP addresses and the timeout is
// 1 minute, the connect to each single address will be given 15 seconds to
// complete before trying the next one.
func (d *Dialer) DialContext(ctx context.Context, network string, address string) (*Conn, error) {
	return d.connect(
		ctx,
		network,
		address,
		ConnConfig{
			ClientID:        d.ClientID,
			TransactionalID: d.TransactionalID,
		},
	)
}

// DialLeader opens a connection to the leader of the partition for a given
// topic.
//
// The address given to the DialContext method may not be the one that the
// connection will end up being established to, because the dialer will lookup
// the partition leader for the topic and return a connection to that server.
// The original address is only used as a mechanism to discover the
// configuration of the kafka cluster that we're connecting to.
func (d *Dialer) DialLeader(ctx context.Context, network string, address string, topic string, partition int) (*Conn, error) {
	p, err := d.LookupPartition(ctx, network, address, topic, partition)
	if err != nil {
		return nil, err
	}
	return d.DialPartition(ctx, network, address, p)
}

// DialPartition opens a connection to the leader of the partition specified by partition
// descriptor. It's strongly advised to use descriptor of the partition that comes out of
// functions LookupPartition or LookupPartitions.
func (d *Dialer) DialPartition(ctx context.Context, network string, address string, partition Partition) (*Conn, error) {
	return d.connect(ctx, network, net.JoinHostPort(partition.Leader.Host, strconv.Itoa(partition.Leader.Port)), ConnConfig{
		ClientID:        d.ClientID,
		Topic:           partition.Topic,
		Partition:       partition.ID,
		TransactionalID: d.TransactionalID,
	})
}

// LookupLeader searches for the kafka broker that is the leader of the
// partition for a given topic, returning a Broker value representing it.
func (d *Dialer) LookupLeader(ctx context.Context, network string, address string, topic string, partition int) (Broker, error) {
	p, err := d.LookupPartition(ctx, network, address, topic, partition)
	return p.Leader, err
}

// LookupPartition searches for the description of specified partition id.
func (d *Dialer) LookupPartition(ctx context.Context, network string, address string, topic string, partition int) (Partition, error) {
	c, err := d.DialContext(ctx, network, address)
	if err != nil {
		return Partition{}, err
	}
	defer c.Close()

	brkch := make(chan Partition, 1)
	errch := make(chan error, 1)

	go func() {
		for attempt := 0; true; attempt++ {
			if attempt != 0 {
				if !sleep(ctx, backoff(attempt, 100*time.Millisecond, 10*time.Second)) {
					errch <- ctx.Err()
					return
				}
			}

			partitions, err := c.ReadPartitions(topic)
			if err != nil {
				if isTemporary(err) {
					continue
				}
				errch <- err
				return
			}

			for _, p := range partitions {
				if p.ID == partition {
					brkch <- p
					return
				}
			}
		}

		errch <- UnknownTopicOrPartition
	}()

	var prt Partition
	select {
	case prt = <-brkch:
	case err = <-errch:
	case <-ctx.Done():
		err = ctx.Err()
	}
	return prt, err
}

// LookupPartitions returns the list of partitions that exist for the given topic.
func (d *Dialer) LookupPartitions(ctx context.Context, network string, address string, topic string) ([]Partition, error) {
	conn, err := d.DialContext(ctx, network, address)
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	prtch := make(chan []Partition, 1)
	errch := make(chan error, 1)

	go func() {
		if prt, err := conn.ReadPartitions(topic); err != nil {
			errch <- err
		} else {
			prtch <- prt
		}
	}()

	var prt []Partition
	select {
	case prt = <-prtch:
	case err = <-errch:
	case <-ctx.Done():
		err = ctx.Err()
	}
	return prt, err
}

// connectTLS returns a tls.Conn that has already completed the Handshake
func (d *Dialer) connectTLS(ctx context.Context, conn net.Conn, config *tls.Config) (tlsConn *tls.Conn, err error) {
	tlsConn = tls.Client(conn, config)
	errch := make(chan error)

	go func() {
		defer close(errch)
		errch <- tlsConn.Handshake()
	}()

	select {
	case <-ctx.Done():
		conn.Close()
		tlsConn.Close()
		<-errch // ignore possible error from Handshake
		err = ctx.Err()

	case err = <-errch:
	}

	return
}

// connect opens a socket connection to the broker, wraps it to create a
// kafka connection, and performs SASL authentication if configured to do so.
func (d *Dialer) connect(ctx context.Context, network, address string, connCfg ConnConfig) (*Conn, error) {
	if d.Timeout != 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, d.Timeout)
		defer cancel()
	}

	if !d.Deadline.IsZero() {
		var cancel context.CancelFunc
		ctx, cancel = context.WithDeadline(ctx, d.Deadline)
		defer cancel()
	}

	c, err := d.dialContext(ctx, network, address)
	if err != nil {
		return nil, err
	}

	conn := NewConnWith(c, connCfg)

	if d.SASLMechanism != nil {
		if err := d.authenticateSASL(ctx, conn); err != nil {
			_ = conn.Close()
			return nil, err
		}
	}

	return conn, nil
}

// authenticateSASL performs all of the required requests to authenticate this
// connection.  If any step fails, this function returns with an error.  A nil
// error indicates successful authentication.
//
// In case of error, this function *does not* close the connection.  That is the
// responsibility of the caller.
func (d *Dialer) authenticateSASL(ctx context.Context, conn *Conn) error {
	if err := conn.saslHandshake(d.SASLMechanism.Name()); err != nil {
		return err
	}

	sess, state, err := d.SASLMechanism.Start(ctx)
	if err != nil {
		return err
	}

	for completed := false; !completed; {
		challenge, err := conn.saslAuthenticate(state)
		switch err {
		case nil:
		case io.EOF:
			// the broker may communicate a failed exchange by closing the
			// connection (esp. in the case where we're passing opaque sasl
			// data over the wire since there's no protocol info).
			return SASLAuthenticationFailed
		default:
			return err
		}

		completed, state, err = sess.Next(ctx, challenge)
		if err != nil {
			return err
		}
	}

	return nil
}

func (d *Dialer) dialContext(ctx context.Context, network string, address string) (net.Conn, error) {
	if r := d.Resolver; r != nil {
		host, port := splitHostPort(address)
		addrs, err := r.LookupHost(ctx, host)
		if err != nil {
			return nil, err
		}
		if len(addrs) != 0 {
			address = addrs[0]
		}
		if len(port) != 0 {
			address, _ = splitHostPort(address)
			address = net.JoinHostPort(address, port)
		}
	}

	conn, err := (&net.Dialer{
		LocalAddr:     d.LocalAddr,
		DualStack:     d.DualStack,
		FallbackDelay: d.FallbackDelay,
		KeepAlive:     d.KeepAlive,
	}).DialContext(ctx, network, address)
	if err != nil {
		return nil, err
	}

	if d.TLS != nil {
		c := d.TLS
		// If no ServerName is set, infer the ServerName
		// from the hostname we're connecting to.
		if c.ServerName == "" {
			c = d.TLS.Clone()
			// Copied from tls.go in the standard library.
			colonPos := strings.LastIndex(address, ":")
			if colonPos == -1 {
				colonPos = len(address)
			}
			hostname := address[:colonPos]
			c.ServerName = hostname
		}
		return d.connectTLS(ctx, conn, c)
	}

	return conn, nil
}

// DefaultDialer is the default dialer used when none is specified.
var DefaultDialer = &Dialer{
	Timeout:   10 * time.Second,
	DualStack: true,
}

// Dial is a convenience wrapper for DefaultDialer.Dial.
func Dial(network string, address string) (*Conn, error) {
	return DefaultDialer.Dial(network, address)
}

// DialContext is a convenience wrapper for DefaultDialer.DialContext.
func DialContext(ctx context.Context, network string, address string) (*Conn, error) {
	return DefaultDialer.DialContext(ctx, network, address)
}

// DialLeader is a convenience wrapper for DefaultDialer.DialLeader.
func DialLeader(ctx context.Context, network string, address string, topic string, partition int) (*Conn, error) {
	return DefaultDialer.DialLeader(ctx, network, address, topic, partition)
}

// DialPartition is a convenience wrapper for DefaultDialer.DialPartition.
func DialPartition(ctx context.Context, network string, address string, partition Partition) (*Conn, error) {
	return DefaultDialer.DialPartition(ctx, network, address, partition)
}

// LookupPartition is a convenience wrapper for DefaultDialer.LookupPartition.
func LookupPartition(ctx context.Context, network string, address string, topic string, partition int) (Partition, error) {
	return DefaultDialer.LookupPartition(ctx, network, address, topic, partition)
}

// LookupPartitions is a convenience wrapper for DefaultDialer.LookupPartitions.
func LookupPartitions(ctx context.Context, network string, address string, topic string) ([]Partition, error) {
	return DefaultDialer.LookupPartitions(ctx, network, address, topic)
}

// The Resolver interface is used as an abstraction to provide service discovery
// of the hosts of a kafka cluster.
type Resolver interface {
	// LookupHost looks up the given host using the local resolver.
	// It returns a slice of that host's addresses.
	LookupHost(ctx context.Context, host string) (addrs []string, err error)
}

func sleep(ctx context.Context, duration time.Duration) bool {
	if duration == 0 {
		select {
		default:
			return true
		case <-ctx.Done():
			return false
		}
	}
	timer := time.NewTimer(duration)
	defer timer.Stop()
	select {
	case <-timer.C:
		return true
	case <-ctx.Done():
		return false
	}
}

func backoff(attempt int, min time.Duration, max time.Duration) time.Duration {
	d := time.Duration(attempt*attempt) * min
	if d > max {
		d = max
	}
	return d
}

func splitHostPort(s string) (host string, port string) {
	host, port, _ = net.SplitHostPort(s)
	if len(host) == 0 && len(port) == 0 {
		host = s
	}
	return
}
<commit_msg>make the dial function configurable on the kafka.Dialer type (#481)<commit_after>package kafka

import (
	"context"
	"crypto/tls"
	"io"
	"net"
	"strconv"
	"strings"
	"time"

	"github.com/segmentio/kafka-go/sasl"
)

// The Dialer type mirrors the net.Dialer API but is designed to open kafka
// connections instead of raw network connections.
type Dialer struct {
	// Unique identifier for client connections established by this Dialer.
	ClientID string

	// Optionally specifies the function that the dialer uses to establish
	// network connections. If nil, net.(*Dialer).DialContext is used instead.
	//
	// When DialFunc is set, LocalAddr, DualStack, FallbackDelay, and KeepAlive
	// are ignored.
	DialFunc func(ctx context.Context, network string, address string) (net.Conn, error)

	// Timeout is the maximum amount of time a dial will wait for a connect to
	// complete. If Deadline is also set, it may fail earlier.
	//
	// The default is no timeout.
	//
	// When dialing a name with multiple IP addresses, the timeout may be
	// divided between them.
	//
	// With or without a timeout, the operating system may impose its own
	// earlier timeout. For instance, TCP timeouts are often around 3 minutes.
	Timeout time.Duration

	// Deadline is the absolute point in time after which dials will fail.
	// If Timeout is set, it may fail earlier.
	// Zero means no deadline, or dependent on the operating system as with the
	// Timeout option.
	Deadline time.Time

	// LocalAddr is the local address to use when dialing an address.
	// The address must be of a compatible type for the network being dialed.
	// If nil, a local address is automatically chosen.
	LocalAddr net.Addr

	// DualStack enables RFC 6555-compliant "Happy Eyeballs" dialing when the
	// network is "tcp" and the destination is a host name with both IPv4 and
	// IPv6 addresses. This allows a client to tolerate networks where one
	// address family is silently broken.
	DualStack bool

	// FallbackDelay specifies the length of time to wait before spawning a
	// fallback connection, when DualStack is enabled.
	// If zero, a default delay of 300ms is used.
	FallbackDelay time.Duration

	// KeepAlive specifies the keep-alive period for an active network
	// connection.
	// If zero, keep-alives are not enabled. Network protocols that do not
	// support keep-alives ignore this field.
	KeepAlive time.Duration

	// Resolver optionally specifies an alternate resolver to use.
	Resolver Resolver

	// TLS enables Dialer to open secure connections.  If nil, standard net.Conn
	// will be used.
	TLS *tls.Config

	// SASLMechanism configures the Dialer to use SASL authentication.  If nil,
	// no authentication will be performed.
	SASLMechanism sasl.Mechanism

	// The transactional id to use for transactional delivery. Idempotent
	// deliver should be enabled if transactional id is configured.
	// For more details look at transactional.id description here: http://kafka.apache.org/documentation.html#producerconfigs
	// Empty string means that the connection will be non-transactional.
	TransactionalID string
}

// Dial connects to the address on the named network.
func (d *Dialer) Dial(network string, address string) (*Conn, error) {
	return d.DialContext(context.Background(), network, address)
}

// DialContext connects to the address on the named network using the provided
// context.
//
// The provided Context must be non-nil. If the context expires before the
// connection is complete, an error is returned. Once successfully connected,
// any expiration of the context will not affect the connection.
//
// When using TCP, and the host in the address parameter resolves to multiple
// network addresses, any dial timeout (from d.Timeout or ctx) is spread over
// each consecutive dial, such that each is given an appropriate fraction of the
// time to connect. For example, if a host has 4 IP addresses and the timeout is
// 1 minute, the connect to each single address will be given 15 seconds to
// complete before trying the next one.
func (d *Dialer) DialContext(ctx context.Context, network string, address string) (*Conn, error) {
	return d.connect(
		ctx,
		network,
		address,
		ConnConfig{
			ClientID:        d.ClientID,
			TransactionalID: d.TransactionalID,
		},
	)
}

// DialLeader opens a connection to the leader of the partition for a given
// topic.
//
// The address given to the DialContext method may not be the one that the
// connection will end up being established to, because the dialer will lookup
// the partition leader for the topic and return a connection to that server.
// The original address is only used as a mechanism to discover the
// configuration of the kafka cluster that we're connecting to.
func (d *Dialer) DialLeader(ctx context.Context, network string, address string, topic string, partition int) (*Conn, error) {
	p, err := d.LookupPartition(ctx, network, address, topic, partition)
	if err != nil {
		return nil, err
	}
	return d.DialPartition(ctx, network, address, p)
}

// DialPartition opens a connection to the leader of the partition specified by partition
// descriptor. It's strongly advised to use descriptor of the partition that comes out of
// functions LookupPartition or LookupPartitions.
func (d *Dialer) DialPartition(ctx context.Context, network string, address string, partition Partition) (*Conn, error) {
	return d.connect(ctx, network, net.JoinHostPort(partition.Leader.Host, strconv.Itoa(partition.Leader.Port)), ConnConfig{
		ClientID:        d.ClientID,
		Topic:           partition.Topic,
		Partition:       partition.ID,
		TransactionalID: d.TransactionalID,
	})
}

// LookupLeader searches for the kafka broker that is the leader of the
// partition for a given topic, returning a Broker value representing it.
func (d *Dialer) LookupLeader(ctx context.Context, network string, address string, topic string, partition int) (Broker, error) {
	p, err := d.LookupPartition(ctx, network, address, topic, partition)
	return p.Leader, err
}

// LookupPartition searches for the description of specified partition id.
func (d *Dialer) LookupPartition(ctx context.Context, network string, address string, topic string, partition int) (Partition, error) {
	c, err := d.DialContext(ctx, network, address)
	if err != nil {
		return Partition{}, err
	}
	defer c.Close()

	brkch := make(chan Partition, 1)
	errch := make(chan error, 1)

	go func() {
		for attempt := 0; true; attempt++ {
			if attempt != 0 {
				if !sleep(ctx, backoff(attempt, 100*time.Millisecond, 10*time.Second)) {
					errch <- ctx.Err()
					return
				}
			}

			partitions, err := c.ReadPartitions(topic)
			if err != nil {
				if isTemporary(err) {
					continue
				}
				errch <- err
				return
			}

			for _, p := range partitions {
				if p.ID == partition {
					brkch <- p
					return
				}
			}
		}

		errch <- UnknownTopicOrPartition
	}()

	var prt Partition
	select {
	case prt = <-brkch:
	case err = <-errch:
	case <-ctx.Done():
		err = ctx.Err()
	}
	return prt, err
}

// LookupPartitions returns the list of partitions that exist for the given topic.
func (d *Dialer) LookupPartitions(ctx context.Context, network string, address string, topic string) ([]Partition, error) {
	conn, err := d.DialContext(ctx, network, address)
	if err != nil {
		return nil, err
	}
	defer conn.Close()

	prtch := make(chan []Partition, 1)
	errch := make(chan error, 1)

	go func() {
		if prt, err := conn.ReadPartitions(topic); err != nil {
			errch <- err
		} else {
			prtch <- prt
		}
	}()

	var prt []Partition
	select {
	case prt = <-prtch:
	case err = <-errch:
	case <-ctx.Done():
		err = ctx.Err()
	}
	return prt, err
}

// connectTLS returns a tls.Conn that has already completed the Handshake
func (d *Dialer) connectTLS(ctx context.Context, conn net.Conn, config *tls.Config) (tlsConn *tls.Conn, err error) {
	tlsConn = tls.Client(conn, config)
	errch := make(chan error)

	go func() {
		defer close(errch)
		errch <- tlsConn.Handshake()
	}()

	select {
	case <-ctx.Done():
		conn.Close()
		tlsConn.Close()
		<-errch // ignore possible error from Handshake
		err = ctx.Err()

	case err = <-errch:
	}

	return
}

// connect opens a socket connection to the broker, wraps it to create a
// kafka connection, and performs SASL authentication if configured to do so.
func (d *Dialer) connect(ctx context.Context, network, address string, connCfg ConnConfig) (*Conn, error) {
	if d.Timeout != 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, d.Timeout)
		defer cancel()
	}

	if !d.Deadline.IsZero() {
		var cancel context.CancelFunc
		ctx, cancel = context.WithDeadline(ctx, d.Deadline)
		defer cancel()
	}

	c, err := d.dialContext(ctx, network, address)
	if err != nil {
		return nil, err
	}

	conn := NewConnWith(c, connCfg)

	if d.SASLMechanism != nil {
		if err := d.authenticateSASL(ctx, conn); err != nil {
			_ = conn.Close()
			return nil, err
		}
	}

	return conn, nil
}

// authenticateSASL performs all of the required requests to authenticate this
// connection.  If any step fails, this function returns with an error.  A nil
// error indicates successful authentication.
//
// In case of error, this function *does not* close the connection.  That is the
// responsibility of the caller.
func (d *Dialer) authenticateSASL(ctx context.Context, conn *Conn) error {
	if err := conn.saslHandshake(d.SASLMechanism.Name()); err != nil {
		return err
	}

	sess, state, err := d.SASLMechanism.Start(ctx)
	if err != nil {
		return err
	}

	for completed := false; !completed; {
		challenge, err := conn.saslAuthenticate(state)
		switch err {
		case nil:
		case io.EOF:
			// the broker may communicate a failed exchange by closing the
			// connection (esp. in the case where we're passing opaque sasl
			// data over the wire since there's no protocol info).
			return SASLAuthenticationFailed
		default:
			return err
		}

		completed, state, err = sess.Next(ctx, challenge)
		if err != nil {
			return err
		}
	}

	return nil
}

func (d *Dialer) dialContext(ctx context.Context, network string, address string) (net.Conn, error) {
	if r := d.Resolver; r != nil {
		host, port := splitHostPort(address)
		addrs, err := r.LookupHost(ctx, host)
		if err != nil {
			return nil, err
		}
		if len(addrs) != 0 {
			address = addrs[0]
		}
		if len(port) != 0 {
			address, _ = splitHostPort(address)
			address = net.JoinHostPort(address, port)
		}
	}

	dial := d.DialFunc
	if dial == nil {
		dial = (&net.Dialer{
			LocalAddr:     d.LocalAddr,
			DualStack:     d.DualStack,
			FallbackDelay: d.FallbackDelay,
			KeepAlive:     d.KeepAlive,
		}).DialContext
	}

	conn, err := dial(ctx, network, address)
	if err != nil {
		return nil, err
	}

	if d.TLS != nil {
		c := d.TLS
		// If no ServerName is set, infer the ServerName
		// from the hostname we're connecting to.
		if c.ServerName == "" {
			c = d.TLS.Clone()
			// Copied from tls.go in the standard library.
			colonPos := strings.LastIndex(address, ":")
			if colonPos == -1 {
				colonPos = len(address)
			}
			hostname := address[:colonPos]
			c.ServerName = hostname
		}
		return d.connectTLS(ctx, conn, c)
	}

	return conn, nil
}

// DefaultDialer is the default dialer used when none is specified.
var DefaultDialer = &Dialer{
	Timeout:   10 * time.Second,
	DualStack: true,
}

// Dial is a convenience wrapper for DefaultDialer.Dial.
func Dial(network string, address string) (*Conn, error) {
	return DefaultDialer.Dial(network, address)
}

// DialContext is a convenience wrapper for DefaultDialer.DialContext.
func DialContext(ctx context.Context, network string, address string) (*Conn, error) {
	return DefaultDialer.DialContext(ctx, network, address)
}

// DialLeader is a convenience wrapper for DefaultDialer.DialLeader.
func DialLeader(ctx context.Context, network string, address string, topic string, partition int) (*Conn, error) {
	return DefaultDialer.DialLeader(ctx, network, address, topic, partition)
}

// DialPartition is a convenience wrapper for DefaultDialer.DialPartition.
func DialPartition(ctx context.Context, network string, address string, partition Partition) (*Conn, error) {
	return DefaultDialer.DialPartition(ctx, network, address, partition)
}

// LookupPartition is a convenience wrapper for DefaultDialer.LookupPartition.
func LookupPartition(ctx context.Context, network string, address string, topic string, partition int) (Partition, error) {
	return DefaultDialer.LookupPartition(ctx, network, address, topic, partition)
}

// LookupPartitions is a convenience wrapper for DefaultDialer.LookupPartitions.
func LookupPartitions(ctx context.Context, network string, address string, topic string) ([]Partition, error) {
	return DefaultDialer.LookupPartitions(ctx, network, address, topic)
}

// The Resolver interface is used as an abstraction to provide service discovery
// of the hosts of a kafka cluster.
type Resolver interface {
	// LookupHost looks up the given host using the local resolver.
	// It returns a slice of that host's addresses.
	LookupHost(ctx context.Context, host string) (addrs []string, err error)
}

func sleep(ctx context.Context, duration time.Duration) bool {
	if duration == 0 {
		select {
		default:
			return true
		case <-ctx.Done():
			return false
		}
	}
	timer := time.NewTimer(duration)
	defer timer.Stop()
	select {
	case <-timer.C:
		return true
	case <-ctx.Done():
		return false
	}
}

func backoff(attempt int, min time.Duration, max time.Duration) time.Duration {
	d := time.Duration(attempt*attempt) * min
	if d > max {
		d = max
	}
	return d
}

func splitHostPort(s string) (host string, port string) {
	host, port, _ = net.SplitHostPort(s)
	if len(host) == 0 && len(port) == 0 {
		host = s
	}
	return
}
<|endoftext|>


<commit_before>package gonduit

import (
	"errors"
	"strings"
)

var (
	// ErrJSONOutputUnsupported is returned when conduit doesn't support JSON output.
	ErrJSONOutputUnsupported = errors.New("json output not supported")

	// ErrURLEncodedInputUnsupported is returned when conduit doesn't support URL encoded input.
	ErrURLEncodedInputUnsupported = errors.New("urlencoded input not supported")

	// ErrSessionAuthUnsupported is returned when conduit doesn't support session authentication.
	ErrSessionAuthUnsupported = errors.New("session authentication not supported")
)

// ConduitError is returned when conduit
// requests return an error response.
type ConduitError struct {
	code string
	info string
}

// Code returns the error_code returned in a conduit response.
func (err *ConduitError) Code() string {
	return err.code
}

// Info returns the error_info returned in a conduit response.
func (err *ConduitError) Info() string {
	return err.info
}

func (err *ConduitError) Error() string {
	return err.code + ": " + err.info
}

// IsConduitError checks whether or not err is a ConduitError.
func IsConduitError(err error) bool {
	_, ok := err.(*ConduitError)
	return ok
}

// A Dialer contains options for connecting to an address.
type Dialer struct {
	ClientName        string
	ClientVersion     string
	ClientDescription string
}

type conduitCapabilitiesResponse struct {
	Authentication []string `json:"authentication"`
	Signatures     []string `json:"signatures"`
	Input          []string `json:"input"`
	Output         []string `json:"output"`
}

// Dial connects to conduit and confirms the API capabilities
// for future calls.
func Dial(host string, options *ClientOptions) (*Conn, error) {
	var d Dialer
	d.ClientName = "gonduit"
	d.ClientVersion = "1"
	return d.Dial(host, options)
}

// Dial connects to conduit and confirms the API capabilities
// for future calls.
func (d *Dialer) Dial(host string, options *ClientOptions) (*Conn, error) {
	host = strings.TrimSuffix(host, "/")

	var resp conduitCapabilitiesResponse
	if err := call(host+"/api/conduit.getcapabilities", nil, &resp, options); err != nil {
		return nil, err
	}

	// We use conduit.connect for authentication
	// and it establishes a session.
	if !containsString(resp.Authentication, "session") {
		return nil, ErrSessionAuthUnsupported
	}

	if !containsString(resp.Input, "urlencoded") {
		return nil, ErrURLEncodedInputUnsupported
	}

	if !containsString(resp.Output, "json") {
		return nil, ErrJSONOutputUnsupported
	}

	conn := Conn{
		host:         host,
		capabilities: &resp,
		dialer:       d,
	}

	return &conn, nil
}
<commit_msg>Store options<commit_after>package gonduit

import (
	"errors"
	"strings"
)

var (
	// ErrJSONOutputUnsupported is returned when conduit doesn't support JSON output.
	ErrJSONOutputUnsupported = errors.New("json output not supported")

	// ErrURLEncodedInputUnsupported is returned when conduit doesn't support URL encoded input.
	ErrURLEncodedInputUnsupported = errors.New("urlencoded input not supported")

	// ErrSessionAuthUnsupported is returned when conduit doesn't support session authentication.
	ErrSessionAuthUnsupported = errors.New("session authentication not supported")
)

// ConduitError is returned when conduit
// requests return an error response.
type ConduitError struct {
	code string
	info string
}

// Code returns the error_code returned in a conduit response.
func (err *ConduitError) Code() string {
	return err.code
}

// Info returns the error_info returned in a conduit response.
func (err *ConduitError) Info() string {
	return err.info
}

func (err *ConduitError) Error() string {
	return err.code + ": " + err.info
}

// IsConduitError checks whether or not err is a ConduitError.
func IsConduitError(err error) bool {
	_, ok := err.(*ConduitError)
	return ok
}

// A Dialer contains options for connecting to an address.
type Dialer struct {
	ClientName        string
	ClientVersion     string
	ClientDescription string
}

type conduitCapabilitiesResponse struct {
	Authentication []string `json:"authentication"`
	Signatures     []string `json:"signatures"`
	Input          []string `json:"input"`
	Output         []string `json:"output"`
}

// Dial connects to conduit and confirms the API capabilities
// for future calls.
func Dial(host string, options *ClientOptions) (*Conn, error) {
	var d Dialer
	d.ClientName = "gonduit"
	d.ClientVersion = "1"
	return d.Dial(host, options)
}

// Dial connects to conduit and confirms the API capabilities
// for future calls.
func (d *Dialer) Dial(host string, options *ClientOptions) (*Conn, error) {
	host = strings.TrimSuffix(host, "/")

	var resp conduitCapabilitiesResponse
	if err := call(host+"/api/conduit.getcapabilities", nil, &resp, options); err != nil {
		return nil, err
	}

	// We use conduit.connect for authentication
	// and it establishes a session.
	if !containsString(resp.Authentication, "session") {
		return nil, ErrSessionAuthUnsupported
	}

	if !containsString(resp.Input, "urlencoded") {
		return nil, ErrURLEncodedInputUnsupported
	}

	if !containsString(resp.Output, "json") {
		return nil, ErrJSONOutputUnsupported
	}

	conn := Conn{
		host:         host,
		capabilities: &resp,
		dialer:       d,
		options:      options,
	}

	return &conn, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2014, Google, Inc.  All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//     http://www.apache.org/licenses/LICENSE-2.0
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// File: rollbackserver.go

package main

import (
	"crypto/tls"
	"crypto/x509"
	"flag"
	"log"
	"net"

	"github.com/jlmucb/cloudproxy/apps/fileproxy"
	"github.com/jlmucb/cloudproxy/tao"
	"github.com/jlmucb/cloudproxy/tao/auth"
	taonet "github.com/jlmucb/cloudproxy/tao/net"
	"github.com/jlmucb/cloudproxy/util"
)

func serve(serverAddr string, prin string, policyCert []byte, signingKey *tao.Keys, policy *fileproxy.ProgramPolicy, m *fileproxy.RollbackMaster) error {
	pc, err := x509.ParseCertificate(policyCert)
	if err != nil {
		return err
	}
	pool := x509.NewCertPool()
	pool.AddCert(pc)
	tlsc, err := taonet.EncodeTLSCert(signingKey)
	if err != nil {
		return err
	}
	conf := &tls.Config{
		RootCAs:            pool,
		Certificates:       []tls.Certificate{*tlsc},
		InsecureSkipVerify: false,
		ClientAuth:         tls.RequireAnyClientCert,
	}
	log.Println("Rollback server listening")
	sock, err := tls.Listen("tcp", serverAddr, conf)
	if err != nil {
		return err
	}

	for {
		conn, err := sock.Accept()
		if err != nil {
			return err
		}
		var clientName string
		if err = conn.(*tls.Conn).Handshake(); err != nil {
			log.Println("TLS handshake failed")
			continue
		}

		peerCerts := conn.(*tls.Conn).ConnectionState().PeerCertificates
		if peerCerts == nil {
			log.Println("rollbackserver: can't get peer list")
			continue
		}

		peerCert := conn.(*tls.Conn).ConnectionState().PeerCertificates[0]
		if peerCert.Raw == nil {
			log.Println("rollbackserver: can't get peer name")
			continue
		}

		if peerCert.Subject.OrganizationalUnit == nil {
			log.Println("No OrganizationalUnit name in the peer certificate. Refusing the connection")
			continue
		}

		clientName = peerCert.Subject.OrganizationalUnit[0]
		ms := util.NewMessageStream(conn)
		// TODO(tmroeder): support multiple simultaneous clients.
		// Add this program as a rollback program.
		log.Printf("Adding a program with name '%s'\n", clientName)
		_ = m.AddRollbackProgram(clientName)
		if err := m.RunMessageLoop(ms, policy, clientName); err != nil {
			log.Printf("rollbackserver: failed to run message loop: %s\n", err)
		}
	}
}

func main() {
	caAddr := flag.String("caAddr", "localhost:8124", "The address of the CA for setting up a certificate signed by the policy key")
	hostcfg := flag.String("hostconfig", "../hostdomain/tao.config", "path to host tao configuration")
	serverHost := flag.String("host", "localhost", "address for client/server")
	serverPort := flag.String("port", "8129", "port for client/server")
	rollbackserverPath := flag.String("rollbackserver_files", "rollbackserver_files/", "rollbackserver directory")

	flag.Parse()
	serverAddr := net.JoinHostPort(*serverHost, *serverPort)

	hostDomain, err := tao.LoadDomain(*hostcfg, nil)
	if err != nil {
		log.Fatalln("rollbackserver: can't load domain")
	}
	var policyCert []byte
	if hostDomain.Keys.Cert != nil {
		policyCert = hostDomain.Keys.Cert.Raw
	}
	if policyCert == nil {
		log.Fatalln("rollbackserver: can't retrieve policy cert")
	}

	if err := hostDomain.ExtendTaoName(tao.Parent()); err != nil {
		log.Fatalln("fileserver: can't extend the Tao with the policy key")
	}
	e := auth.PrinExt{Name: "rollbackserver_version_1"}
	if err = tao.Parent().ExtendTaoName(auth.SubPrin{e}); err != nil {
		log.Fatalln("rollbackserver: can't extend name")
	}

	taoName, err := tao.Parent().GetTaoName()
	if err != nil {
		return
	}

	sealedSymmetricKey, sealedSigningKey, programCert, delegation, err := fileproxy.LoadProgramKeys(*rollbackserverPath)
	if err != nil {
		log.Println("rollbackserver: can't retrieve key material")
	}
	if sealedSymmetricKey == nil || sealedSigningKey == nil || delegation == nil || programCert == nil {
		log.Println("rollbackserver: No key material present")
	}

	var symKeys []byte
	defer fileproxy.ZeroBytes(symKeys)
	if sealedSymmetricKey != nil {
		var policy string
		if symKeys, policy, err = tao.Parent().Unseal(sealedSymmetricKey); err != nil {
			log.Fatalln("rollbackserver: couldn't unseal the symmetric key")
		}
		if policy != tao.SealPolicyDefault {
			log.Fatalln("rollbackserver: unexpected policy on unseal")
		}
	} else {
		if symKeys, err = fileproxy.InitializeSealedSymmetricKeys(*rollbackserverPath, tao.Parent(), fileproxy.SymmetricKeySize); err != nil {
			log.Fatalf("rollbackserver: InitializeSealedSymmetricKeys error: %s\n", err)
		}
	}

	var signingKey *tao.Keys
	if sealedSigningKey != nil {
		if signingKey, err = fileproxy.SigningKeyFromBlob(tao.Parent(), sealedSigningKey, programCert, delegation); err != nil {
			log.Fatalf("rollbackserver: SigningKeyFromBlob error: %s\n", err)
		}
	} else {
		if signingKey, err = fileproxy.InitializeSealedSigningKey(*caAddr, *rollbackserverPath, tao.Parent(), *hostDomain); err != nil {
			log.Fatalf("rollbackserver: InitializeSealedSigningKey error: %s\n", err)
		}
		programCert = signingKey.Cert.Raw
	}

	progPolicy := fileproxy.NewProgramPolicy(policyCert, taoName.String(), signingKey, symKeys, programCert)
	m := fileproxy.NewRollbackMaster(taoName.String())

	if err := serve(serverAddr, taoName.String(), policyCert, signingKey, progPolicy, m); err != nil {
		log.Fatalf("rollbackserver: server error: %s\n", err)
	}
	log.Println("rollbackserver: done")
}
<commit_msg>Port the rollback server to use keys functions.<commit_after>// Copyright (c) 2014, Google, Inc.  All rights reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//     http://www.apache.org/licenses/LICENSE-2.0
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// File: rollbackserver.go

package main

import (
	"crypto/tls"
	"crypto/x509"
	"flag"
	"log"
	"net"

	"github.com/jlmucb/cloudproxy/apps/fileproxy"
	"github.com/jlmucb/cloudproxy/tao"
	"github.com/jlmucb/cloudproxy/tao/auth"
	taonet "github.com/jlmucb/cloudproxy/tao/net"
	"github.com/jlmucb/cloudproxy/util"
)

func serve(serverAddr string, prin string, policyCert []byte, signingKey *tao.Keys, policy *fileproxy.ProgramPolicy, m *fileproxy.RollbackMaster) error {
	pc, err := x509.ParseCertificate(policyCert)
	if err != nil {
		return err
	}
	pool := x509.NewCertPool()
	pool.AddCert(pc)
	tlsc, err := taonet.EncodeTLSCert(signingKey)
	if err != nil {
		return err
	}
	conf := &tls.Config{
		RootCAs:            pool,
		Certificates:       []tls.Certificate{*tlsc},
		InsecureSkipVerify: false,
		ClientAuth:         tls.RequireAnyClientCert,
	}
	log.Println("Rollback server listening")
	sock, err := tls.Listen("tcp", serverAddr, conf)
	if err != nil {
		return err
	}

	for {
		conn, err := sock.Accept()
		if err != nil {
			return err
		}
		var clientName string
		if err = conn.(*tls.Conn).Handshake(); err != nil {
			log.Println("TLS handshake failed")
			continue
		}

		peerCerts := conn.(*tls.Conn).ConnectionState().PeerCertificates
		if peerCerts == nil {
			log.Println("rollbackserver: can't get peer list")
			continue
		}

		peerCert := conn.(*tls.Conn).ConnectionState().PeerCertificates[0]
		if peerCert.Raw == nil {
			log.Println("rollbackserver: can't get peer name")
			continue
		}

		if peerCert.Subject.OrganizationalUnit == nil {
			log.Println("No OrganizationalUnit name in the peer certificate. Refusing the connection")
			continue
		}

		clientName = peerCert.Subject.OrganizationalUnit[0]
		ms := util.NewMessageStream(conn)
		// TODO(tmroeder): support multiple simultaneous clients.
		// Add this program as a rollback program.
		log.Printf("Adding a program with name '%s'\n", clientName)
		_ = m.AddRollbackProgram(clientName)
		if err := m.RunMessageLoop(ms, policy, clientName); err != nil {
			log.Printf("rollbackserver: failed to run message loop: %s\n", err)
		}
	}
}

func main() {
	caAddr := flag.String("caAddr", "localhost:8124", "The address of the CA for setting up a certificate signed by the policy key")
	hostcfg := flag.String("hostconfig", "tao.config", "path to host tao configuration")
	serverHost := flag.String("host", "localhost", "address for client/server")
	serverPort := flag.String("port", "8129", "port for client/server")
	rollbackServerPath := flag.String("rollbackserver_files", "rollbackserver_files", "rollbackserver directory")
	country := flag.String("country", "US", "The country for the fileclient certificate")
	org := flag.String("organization", "Google", "The organization for the fileclient certificate")

	flag.Parse()
	serverAddr := net.JoinHostPort(*serverHost, *serverPort)

	hostDomain, err := tao.LoadDomain(*hostcfg, nil)
	if err != nil {
		log.Fatalln("rollbackserver: can't load domain:", err)
	}
	var policyCert []byte
	if hostDomain.Keys.Cert != nil {
		policyCert = hostDomain.Keys.Cert.Raw
	}
	if policyCert == nil {
		log.Fatalln("rollbackserver: can't retrieve policy cert")
	}

	parentTao := tao.Parent()
	if err := hostDomain.ExtendTaoName(parentTao); err != nil {
		log.Fatalln("fileserver: can't extend the Tao with the policy key")
	}
	e := auth.PrinExt{Name: "rollbackserver_version_1"}
	if err = parentTao.ExtendTaoName(auth.SubPrin{e}); err != nil {
		log.Fatalln("rollbackserver: can't extend name")
	}

	taoName, err := parentTao.GetTaoName()
	if err != nil {
		return
	}

	// Create or read the keys for rollbackserver.
	rbKeys, err := tao.NewOnDiskTaoSealedKeys(tao.Signing|tao.Crypting, parentTao, *rollbackServerPath, tao.SealPolicyDefault)
	if err != nil {
		log.Fatalln("rollbackserver: couldn't set up the Tao-sealed keys:", err)
	}

	// Set up a temporary cert for communication with keyNegoServer.
	rbKeys.Cert, err = rbKeys.SigningKey.CreateSelfSignedX509(tao.NewX509Name(tao.X509Details{
		Country:      *country,
		Organization: *org,
		CommonName:   taoName.String(),
	}))
	if err != nil {
		log.Fatalln("rollbackserver: couldn't create a self-signed cert for rollbackserver keys:", err)
	}

	// Contact keyNegoServer for the certificate.
	if err := fileproxy.EstablishCert("tcp", *caAddr, rbKeys, hostDomain.Keys.VerifyingKey); err != nil {
		log.Fatalf("rollbackserver: couldn't establish a cert signed by the policy key: %s", err)
	}

	// The symmetric keys aren't used by the rollback server.
	progPolicy := fileproxy.NewProgramPolicy(policyCert, taoName.String(), rbKeys, nil, rbKeys.Cert.Raw)
	m := fileproxy.NewRollbackMaster(taoName.String())

	if err := serve(serverAddr, taoName.String(), policyCert, rbKeys, progPolicy, m); err != nil {
		log.Fatalf("rollbackserver: server error: %s\n", err)
	}
	log.Println("rollbackserver: done")
}
<|endoftext|>


<commit_before>package collector

import (
	"github.com/aws/aws-sdk-go/service/elb"
	"github.com/prometheus/client_golang/prometheus"
	"golang.org/x/sync/errgroup"

	"github.com/giantswarm/microerror"
	"github.com/giantswarm/micrologger"

	clientaws "github.com/giantswarm/aws-operator/client/aws"
)

const (
	labelELB = "elb"
)

const (
	subsystemELB = "elb"
)

const (
	stateOutOfService = "OutOfService"
)

var (
	elbsDesc *prometheus.Desc = prometheus.NewDesc(
		prometheus.BuildFQName(namespace, subsystemELB, "instance_out_of_service_count"),
		"Gauge about ELB instances being out of service.",
		[]string{
			labelELB,
			labelAccount,
			labelCluster,
			labelInstallation,
			labelOrganization,
		},
		nil,
	)
)

type ELBConfig struct {
	Helper *helper
	Logger micrologger.Logger

	InstallationName string
}

type ELB struct {
	helper *helper
	logger micrologger.Logger

	installationName string
}

type loadBalancer struct {
	InstancesOutOfService float64
	Name                  string
	Tags                  map[string]string
}

func NewELB(config ELBConfig) (*ELB, error) {
	if config.Helper == nil {
		return nil, microerror.Maskf(invalidConfigError, "%T.Helper must not be empty", config)
	}
	if config.Logger == nil {
		return nil, microerror.Maskf(invalidConfigError, "%T.Logger must not be empty", config)
	}

	if config.InstallationName == "" {
		return nil, microerror.Maskf(invalidConfigError, "%T.InstallationName must not be empty", config)
	}

	e := &ELB{
		helper: config.Helper,
		logger: config.Logger,

		installationName: config.InstallationName,
	}

	return e, nil
}

func (e *ELB) Collect(ch chan<- prometheus.Metric) error {
	awsClientsList, err := e.helper.GetAWSClients()
	if err != nil {
		return microerror.Mask(err)
	}

	var g errgroup.Group

	for _, item := range awsClientsList {
		awsClients := item

		g.Go(func() error {
			err := e.collectForAccount(ch, awsClients)
			if err != nil {
				return microerror.Mask(err)
			}

			return nil
		})
	}

	err = g.Wait()
	if err != nil {
		return microerror.Mask(err)
	}

	return nil
}

func (e *ELB) Describe(ch chan<- *prometheus.Desc) error {
	ch <- elbsDesc
	return nil
}

func (e *ELB) collectForAccount(ch chan<- prometheus.Metric, awsClients clientaws.Clients) error {
	account, err := e.helper.AWSAccountID(awsClients)
	if err != nil {
		return microerror.Mask(err)
	}

	var loadbalancers []*elb.LoadBalancerDescription
	{
		i := &elb.DescribeLoadBalancersInput{}
		o, err := awsClients.ELB.DescribeLoadBalancers(i)
		if err != nil {
			return microerror.Mask(err)
		}
		loadbalancers = o.LoadBalancerDescriptions
	}

	var lbs []loadBalancer
	{
		i := &elb.DescribeTagsInput{}
		for _, l := range loadbalancers {
			i.LoadBalancerNames = append(i.LoadBalancerNames, l.LoadBalancerName)
		}

		o, err := awsClients.ELB.DescribeTags(i)
		if err != nil {
			return microerror.Mask(err)
		}

		for _, d := range o.TagDescriptions {
			lb := loadBalancer{
				Name: *d.LoadBalancerName,
				Tags: make(map[string]string),
			}

			for _, t := range d.Tags {
				lb.Tags[*t.Key] = *t.Value
			}

			if lb.Tags[tagInstallation] != e.installationName {
				continue
			}

			lbs = append(lbs, lb)
		}
	}

	{
		// AWS API doesn't provide a method to describe instance health for all
		// specified ELBs so it must be done with N API calls.
		for _, lb := range lbs {
			i := &elb.DescribeInstanceHealthInput{
				LoadBalancerName: &lb.Name,
			}

			o, err := awsClients.ELB.DescribeInstanceHealth(i)
			if err != nil {
				return microerror.Mask(err)
			}

			for _, s := range o.InstanceStates {
				if *s.State == stateOutOfService {
					lb.InstancesOutOfService++
				}
			}
		}
	}

	{
		for _, lb := range lbs {
			ch <- prometheus.MustNewConstMetric(
				elbsDesc,
				prometheus.GaugeValue,
				lb.InstancesOutOfService,
				lb.Name,
				account,
				lb.Tags[tagCluster],
				lb.Tags[tagInstallation],
				lb.Tags[tagOrganization],
			)
		}
	}

	return nil
}
<commit_msg>Short circuit ELB metrics collection when no LBs present (#1324)<commit_after>package collector

import (
	"github.com/aws/aws-sdk-go/service/elb"
	"github.com/prometheus/client_golang/prometheus"
	"golang.org/x/sync/errgroup"

	"github.com/giantswarm/microerror"
	"github.com/giantswarm/micrologger"

	clientaws "github.com/giantswarm/aws-operator/client/aws"
)

const (
	labelELB = "elb"
)

const (
	subsystemELB = "elb"
)

const (
	stateOutOfService = "OutOfService"
)

var (
	elbsDesc *prometheus.Desc = prometheus.NewDesc(
		prometheus.BuildFQName(namespace, subsystemELB, "instance_out_of_service_count"),
		"Gauge about ELB instances being out of service.",
		[]string{
			labelELB,
			labelAccount,
			labelCluster,
			labelInstallation,
			labelOrganization,
		},
		nil,
	)
)

type ELBConfig struct {
	Helper *helper
	Logger micrologger.Logger

	InstallationName string
}

type ELB struct {
	helper *helper
	logger micrologger.Logger

	installationName string
}

type loadBalancer struct {
	InstancesOutOfService float64
	Name                  string
	Tags                  map[string]string
}

func NewELB(config ELBConfig) (*ELB, error) {
	if config.Helper == nil {
		return nil, microerror.Maskf(invalidConfigError, "%T.Helper must not be empty", config)
	}
	if config.Logger == nil {
		return nil, microerror.Maskf(invalidConfigError, "%T.Logger must not be empty", config)
	}

	if config.InstallationName == "" {
		return nil, microerror.Maskf(invalidConfigError, "%T.InstallationName must not be empty", config)
	}

	e := &ELB{
		helper: config.Helper,
		logger: config.Logger,

		installationName: config.InstallationName,
	}

	return e, nil
}

func (e *ELB) Collect(ch chan<- prometheus.Metric) error {
	awsClientsList, err := e.helper.GetAWSClients()
	if err != nil {
		return microerror.Mask(err)
	}

	var g errgroup.Group

	for _, item := range awsClientsList {
		awsClients := item

		g.Go(func() error {
			err := e.collectForAccount(ch, awsClients)
			if err != nil {
				return microerror.Mask(err)
			}

			return nil
		})
	}

	err = g.Wait()
	if err != nil {
		return microerror.Mask(err)
	}

	return nil
}

func (e *ELB) Describe(ch chan<- *prometheus.Desc) error {
	ch <- elbsDesc
	return nil
}

func (e *ELB) collectForAccount(ch chan<- prometheus.Metric, awsClients clientaws.Clients) error {
	account, err := e.helper.AWSAccountID(awsClients)
	if err != nil {
		return microerror.Mask(err)
	}

	var loadBalancers []*elb.LoadBalancerDescription
	{
		i := &elb.DescribeLoadBalancersInput{}
		o, err := awsClients.ELB.DescribeLoadBalancers(i)
		if err != nil {
			return microerror.Mask(err)
		}
		loadBalancers = o.LoadBalancerDescriptions

		if len(loadBalancers) == 0 {
			// E.g. during cluster creation there are no load balancers present
			// yet so further AWS API calls would fail on validation. No
			// metrics to emit either so we can short circuit here.
			return nil
		}
	}

	var lbs []loadBalancer
	{
		i := &elb.DescribeTagsInput{}
		for _, l := range loadBalancers {
			i.LoadBalancerNames = append(i.LoadBalancerNames, l.LoadBalancerName)
		}

		o, err := awsClients.ELB.DescribeTags(i)
		if err != nil {
			return microerror.Mask(err)
		}

		for _, d := range o.TagDescriptions {
			lb := loadBalancer{
				Name: *d.LoadBalancerName,
				Tags: make(map[string]string),
			}

			for _, t := range d.Tags {
				lb.Tags[*t.Key] = *t.Value
			}

			if lb.Tags[tagInstallation] != e.installationName {
				continue
			}

			lbs = append(lbs, lb)
		}
	}

	{
		// AWS API doesn't provide a method to describe instance health for all
		// specified ELBs so it must be done with N API calls.
		for _, lb := range lbs {
			i := &elb.DescribeInstanceHealthInput{
				LoadBalancerName: &lb.Name,
			}

			o, err := awsClients.ELB.DescribeInstanceHealth(i)
			if err != nil {
				return microerror.Mask(err)
			}

			for _, s := range o.InstanceStates {
				if *s.State == stateOutOfService {
					lb.InstancesOutOfService++
				}
			}
		}
	}

	{
		for _, lb := range lbs {
			ch <- prometheus.MustNewConstMetric(
				elbsDesc,
				prometheus.GaugeValue,
				lb.InstancesOutOfService,
				lb.Name,
				account,
				lb.Tags[tagCluster],
				lb.Tags[tagInstallation],
				lb.Tags[tagOrganization],
			)
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package upload

import (
	"time"

	"github.com/emicklei/go-restful"
	"github.com/materials-commons/mcstore/pkg/db/schema"
	"github.com/materials-commons/mcstore/server/mcstored/service/uploads"
)

// uploadCreateRequest describes the JSON request a client will send
// to create a new upload request.
type uploadCreateRequest struct {
	projectID   string `json:"project_id"`
	directoryID string `json:"directory_id"`
	userID      string `json:"user_id"`
}

// uploadCreateResponse is the format of JSON sent back containing
// the upload request ID.
type uploadCreateResponse struct {
	requestID string `json:"request_id"`
}

// createUploadRequest services requests to create a new upload id. It validates
// the given request, and ensures that the returned upload id is unique. Upload
// requests are persisted until deleted or a successful upload occurs.
func (r *uploadResource) createUploadRequest(request *restful.Request, response *restful.Response, user schema.User) (interface{}, error) {
	var req uploadCreateRequest
	if err := request.ReadEntity(&req); err != nil {
		return nil, err
	}

	cr := uploads.CreateRequest{
		User:        req.userID,
		DirectoryID: req.directoryID,
		ProjectID:   req.projectID,
		Host:        request.Request.RemoteAddr,
		Birthtime:   time.Now(),
	}
	upload, err := r.createService.Create(cr)
	if err != nil {
		return nil, err
	}

	resp := uploadCreateResponse{
		requestID: upload.ID,
	}
	return &resp, nil
}
<commit_msg>Make the fields names public for marshalling and unmarshalling. Add fields for the file.<commit_after>package upload

import (
	"time"

	"github.com/emicklei/go-restful"
	"github.com/materials-commons/mcstore/pkg/db/schema"
	"github.com/materials-commons/mcstore/server/mcstored/service/uploads"
)

// uploadCreateRequest describes the JSON request a client will send
// to create a new upload request.
type uploadCreateRequest struct {
	ProjectID   string `json:"project_id"`
	DirectoryID string `json:"directory_id"`
	FileName    string `json:"filename"`
	FileSize    int64  `json:"filesize"`
	FileCTime   int64  `json:"filectime"`
	UserID      string `json:"user_id"`
}

// uploadCreateResponse is the format of JSON sent back containing
// the upload request ID.
type uploadCreateResponse struct {
	RequestID string `json:"request_id"`
}

// createUploadRequest services requests to create a new upload id. It validates
// the given request, and ensures that the returned upload id is unique. Upload
// requests are persisted until deleted or a successful upload occurs.
func (r *uploadResource) createUploadRequest(request *restful.Request, response *restful.Response, user schema.User) (interface{}, error) {
	var req uploadCreateRequest
	if err := request.ReadEntity(&req); err != nil {
		return nil, err
	}

	cr := uploads.CreateRequest{
		User:        req.UserID,
		DirectoryID: req.DirectoryID,
		ProjectID:   req.ProjectID,
		FileName:    req.FileName,
		FileSize:    req.FileSize,
		FileCTime:   time.Unix(req.FileCTime, 0),
		Host:        request.Request.RemoteAddr,
		Birthtime:   time.Now(),
	}
	upload, err := r.createService.Create(cr)
	if err != nil {
		return nil, err
	}

	resp := uploadCreateResponse{
		RequestID: upload.ID,
	}
	return &resp, nil
}
<|endoftext|>


<commit_before>package easyss

import (
	"bufio"
	"bytes"
	"context"
	"crypto/tls"
	_ "embed"
	"fmt"
	"io"
	"net"
	"net/http"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/coocood/freecache"
	"github.com/miekg/dns"
	"github.com/nange/easypool"
	"github.com/nange/easyss/util"
	"github.com/oschwald/geoip2-golang"
	utls "github.com/refraction-networking/utls"
	log "github.com/sirupsen/logrus"
	"github.com/txthinking/socks5"
)

const version = "v1.5.0"

var (
	//go:embed geodata/geoip_cn_private.mmdb
	geoIPCNPrivate []byte
	//go:embed geodata/geosite_cn.txt
	geoSiteCN []byte
)

func PrintVersion() {
	fmt.Println("easyss version", version)
}

type Statistics struct {
	BytesSend    atomic.Int64
	BytesReceive atomic.Int64
}

type GeoSite struct {
	domain       map[string]struct{}
	fullDomain   map[string]struct{}
	regexpDomain []*regexp.Regexp
}

func NewGeoSite(data []byte) *GeoSite {
	gs := &GeoSite{
		domain:     make(map[string]struct{}),
		fullDomain: make(map[string]struct{}),
	}

	r := bufio.NewReader(bytes.NewReader(data))
	for {
		line, _, err := r.ReadLine()
		if err == io.EOF {
			break
		}

		if bytes.HasPrefix(line, []byte("full:")) {
			gs.fullDomain[string(line[5:])] = struct{}{}
			continue
		}

		if bytes.HasPrefix(line, []byte("regexp:")) {
			line = line[7:]
			re, err := regexp.Compile(string(line))
			if err != nil {
				log.Errorf("compile geosite string:%s, err:%s", string(line), err.Error())
				continue
			}
			gs.regexpDomain = append(gs.regexpDomain, re)
			continue
		}

		gs.domain[string(line)] = struct{}{}
	}

	return gs
}

func (gs *GeoSite) SiteAtCN(domain string) bool {
	domainRoot := func(_domain string) string {
		var firstDot, lastDot int
		for {
			firstDot = strings.Index(_domain, ".")
			lastDot = strings.LastIndex(_domain, ".")
			if firstDot == lastDot {
				return _domain
			}
			_domain = _domain[firstDot+1:]
		}
	}

	if _, ok := gs.fullDomain[domain]; ok {
		return true
	}

	_domain := domainRoot(domain)
	if _, ok := gs.domain[_domain]; ok {
		return true
	}

	for _, re := range gs.regexpDomain {
		if re.MatchString(domain) {
			return true
		}
	}

	return false
}

type Easyss struct {
	config         *Config
	serverIP       string
	stat           *Statistics
	localGw        string
	localDev       string
	devIndex       int
	dnsCache       *freecache.Cache
	directDNSCache *freecache.Cache
	geoipDB        *geoip2.Reader
	geosite        *GeoSite

	// the mu Mutex to protect below fields
	mu              *sync.RWMutex
	tcpPool         easypool.Pool
	socksServer     *socks5.Server
	httpProxyServer *http.Server
	closing         chan struct{}
	tun2socksStatus Tun2socksStatus
}

func New(config *Config) (*Easyss, error) {
	ss := &Easyss{
		config:         config,
		stat:           &Statistics{},
		dnsCache:       freecache.NewCache(1024 * 1024),
		directDNSCache: freecache.NewCache(1024 * 1024),
		closing:        make(chan struct{}, 1),
		mu:             &sync.RWMutex{},
	}

	db, err := geoip2.FromBytes(geoIPCNPrivate)
	if err != nil {
		panic(err)
	}
	ss.geoipDB = db
	ss.geosite = NewGeoSite(geoSiteCN)

	msg, err := ss.ServerDNSMsg()
	if err != nil {
		log.Errorf("query server dns msg err:%s", err.Error())
	}
	if msg != nil {
		ss.serverIP = msg.Answer[0].(*dns.A).A.String()
		ss.SetDNSCache(msg, true, true)
		ss.SetDNSCache(msg, true, false)
	}

	switch runtime.GOOS {
	case "linux", "windows", "darwin":
		gw, dev, err := util.SysGatewayAndDevice()
		if err != nil {
			log.Errorf("get system gateway and device err:%s", err.Error())
		}
		ss.localGw = gw
		ss.localDev = dev

		iface, err := net.InterfaceByName(dev)
		if err != nil {
			log.Errorf("interface by name err:%v", err)
			return nil, err
		}
		ss.devIndex = iface.Index
	}

	go ss.printStatistics()

	return ss, err
}

func (ss *Easyss) InitTcpPool() error {
	if ss.DisableUTLS() {
		log.Infof("uTLS is disabled")
	} else {
		log.Infof("uTLS is enabled")
	}

	factory := func() (net.Conn, error) {
		ctx, cancel := context.WithTimeout(context.Background(), ss.Timeout())
		defer cancel()

		if ss.DisableUTLS() {
			dialer := new(tls.Dialer)
			return dialer.DialContext(ctx, "tcp", ss.ServerAddr())
		}

		conn, err := net.DialTimeout("tcp", ss.ServerAddr(), ss.Timeout())
		if err != nil {
			return nil, err
		}

		uConn := utls.UClient(conn, &utls.Config{ServerName: ss.Server()}, utls.HelloChrome_Auto)
		if err := uConn.HandshakeContext(ctx); err != nil {
			return nil, err
		}

		return uConn, nil
	}

	config := &easypool.PoolConfig{
		InitialCap:  10,
		MaxCap:      50,
		MaxIdle:     10,
		Idletime:    5 * time.Minute,
		MaxLifetime: 30 * time.Minute,
		Factory:     factory,
	}
	tcpPool, err := easypool.NewHeapPool(config)
	ss.SetPool(tcpPool)

	return err
}

func (ss *Easyss) LocalPort() int {
	return ss.config.LocalPort
}

func (ss *Easyss) LocalHttpProxyPort() int {
	return ss.config.LocalPort + 1000
}

func (ss *Easyss) LocalPacPort() int {
	return ss.config.LocalPort + 1001
}

func (ss *Easyss) ServerPort() int {
	return ss.config.ServerPort
}

func (ss *Easyss) Password() string {
	return ss.config.Password
}

func (ss *Easyss) Method() string {
	return ss.config.Method
}

func (ss *Easyss) Server() string {
	return ss.config.Server
}

func (ss *Easyss) ServerIP() string {
	return ss.serverIP
}

func (ss *Easyss) ServerAddr() string {
	return fmt.Sprintf("%s:%d", ss.Server(), ss.ServerPort())
}

func (ss *Easyss) Socks5ProxyAddr() string {
	return fmt.Sprintf("socks5://%s", ss.LocalAddr())
}

func (ss *Easyss) LocalGateway() string {
	return ss.localGw
}

func (ss *Easyss) LocalDevice() string {
	return ss.localDev
}

func (ss *Easyss) LocalDeviceIndex() int {
	return ss.devIndex
}

func (ss *Easyss) Timeout() time.Duration {
	return time.Duration(ss.config.Timeout) * time.Second
}

func (ss *Easyss) LocalAddr() string {
	return fmt.Sprintf("%s:%d", "127.0.0.1", ss.LocalPort())
}

func (ss *Easyss) BindAll() bool {
	return ss.config.BindALL
}

func (ss *Easyss) DisableUTLS() bool {
	return ss.config.DisableUTLS
}

func (ss *Easyss) ConfigFilename() string {
	if ss.config.ConfigFile == "" {
		return ""
	}
	return filepath.Base(ss.config.ConfigFile)
}

func (ss *Easyss) Pool() easypool.Pool {
	ss.mu.RLock()
	defer ss.mu.RUnlock()
	return ss.tcpPool
}

func (ss *Easyss) SetSocksServer(server *socks5.Server) {
	ss.mu.Lock()
	defer ss.mu.Unlock()
	ss.socksServer = server
}

func (ss *Easyss) Tun2socksStatus() Tun2socksStatus {
	ss.mu.RLock()
	defer ss.mu.RUnlock()
	return ss.tun2socksStatus
}

func (ss *Easyss) SetTun2socksStatus(status Tun2socksStatus) {
	ss.mu.Lock()
	defer ss.mu.Unlock()
	ss.tun2socksStatus = status
}

func (ss *Easyss) Tun2socksStatusAuto() bool {
	ss.mu.RLock()
	defer ss.mu.RUnlock()
	return ss.tun2socksStatus == Tun2socksStatusAuto
}

func (ss *Easyss) Tun2socksStatusOn() bool {
	ss.mu.RLock()
	defer ss.mu.RUnlock()
	return ss.tun2socksStatus == Tun2socksStatusOn
}

func (ss *Easyss) Tun2socksStatusOff() bool {
	ss.mu.RLock()
	defer ss.mu.RUnlock()
	return ss.tun2socksStatus == Tun2socksStatusOff
}

func (ss *Easyss) SetHttpProxyServer(server *http.Server) {
	ss.mu.Lock()
	defer ss.mu.Unlock()
	ss.httpProxyServer = server
}

func (ss *Easyss) SetPool(pool easypool.Pool) {
	ss.mu.Lock()
	defer ss.mu.Unlock()
	ss.tcpPool = pool
}

func (ss *Easyss) DNSCache(name, qtype string, isDirect bool) *dns.Msg {
	var v []byte
	var err error
	if isDirect {
		v, err = ss.directDNSCache.Get([]byte(name + qtype))
	} else {
		v, err = ss.dnsCache.Get([]byte(name + qtype))
	}
	if err != nil || len(v) == 0 {
		return nil
	}

	msg := &dns.Msg{}
	if err := msg.Unpack(v); err != nil {
		return nil
	}

	return msg
}

func (ss *Easyss) RenewDNSCache(name, qtype string, isDirect bool) {
	if isDirect {
		ss.directDNSCache.Touch([]byte(name+qtype), 8*60*60)
		return
	}
	ss.dnsCache.Touch([]byte(name+qtype), 8*60*60)
}

func (ss *Easyss) SetDNSCache(msg *dns.Msg, noExpire, isDirect bool) error {
	if msg == nil {
		return nil
	}
	if len(msg.Question) == 0 {
		return nil
	}

	q := msg.Question[0]
	if q.Qtype == dns.TypeA || q.Qtype == dns.TypeAAAA {
		v, err := msg.Pack()
		if err != nil {
			return err
		}
		expireSec := 8 * 60 * 60
		if noExpire {
			expireSec = 0
		}
		key := []byte(q.Name + dns.TypeToString[q.Qtype])
		if isDirect {
			return ss.directDNSCache.Set(key, v, expireSec)
		}
		return ss.dnsCache.Set(key, v, expireSec)
	}

	return nil
}

func (ss *Easyss) ServerDNSMsg() (*dns.Msg, error) {
	c := new(dns.Client)

	m := new(dns.Msg)
	m.SetQuestion(dns.Fqdn(ss.Server()), dns.TypeA)
	m.RecursionDesired = true

	r, _, err := c.Exchange(m, "114.114.114.114:53")
	if err != nil {
		return nil, err
	}

	if r.Rcode != dns.RcodeSuccess {
		return nil, fmt.Errorf("dns query response Rcode:%v not equals RcodeSuccess", r.Rcode)
	}

	return r, nil
}

func (ss *Easyss) HostAtCN(host string) bool {
	if host == "" {
		return false
	}

	if util.IsIP(host) {
		return ss.IPAtCN(host)
	}

	return ss.geosite.SiteAtCN(host)
}

func (ss *Easyss) IPAtCN(ip string) bool {
	_ip := net.ParseIP(ip)
	if _ip == nil {
		return false
	}
	country, err := ss.geoipDB.Country(_ip)
	if err != nil {
		return false
	}

	if country.Country.IsoCode == "CN" {
		return true
	}

	return false
}

func (ss *Easyss) Close() {
	ss.mu.Lock()
	defer ss.mu.Unlock()

	if ss.tcpPool != nil {
		ss.tcpPool.Close()
		ss.tcpPool = nil
	}
	if ss.httpProxyServer != nil {
		ss.httpProxyServer.Close()
		ss.httpProxyServer = nil
	}
	if ss.socksServer != nil {
		ss.socksServer.Shutdown()
		ss.socksServer = nil
	}
	if ss.closing != nil {
		close(ss.closing)
		ss.closing = nil
	}
	if ss.tun2socksStatus != Tun2socksStatusOff {
		ss.closeTun2socks()
	}
}

func (ss *Easyss) printStatistics() {
	ss.mu.Lock()
	closing := ss.closing
	ss.mu.Unlock()
	for {
		select {
		case <-time.After(time.Hour):
			sendSize := ss.stat.BytesSend.Load() / (1024 * 1024)
			receiveSize := ss.stat.BytesReceive.Load() / (1024 * 1024)
			log.Debugf("easyss send data size: %vMB, recive data size: %vMB", sendSize, receiveSize)
		case <-closing:
			return
		}
	}
}
<commit_msg>easyss: close tun2socks service first if it is on<commit_after>package easyss

import (
	"bufio"
	"bytes"
	"context"
	"crypto/tls"
	_ "embed"
	"fmt"
	"io"
	"net"
	"net/http"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/coocood/freecache"
	"github.com/miekg/dns"
	"github.com/nange/easypool"
	"github.com/nange/easyss/util"
	"github.com/oschwald/geoip2-golang"
	utls "github.com/refraction-networking/utls"
	log "github.com/sirupsen/logrus"
	"github.com/txthinking/socks5"
)

const version = "v1.5.0"

var (
	//go:embed geodata/geoip_cn_private.mmdb
	geoIPCNPrivate []byte
	//go:embed geodata/geosite_cn.txt
	geoSiteCN []byte
)

func PrintVersion() {
	fmt.Println("easyss version", version)
}

type Statistics struct {
	BytesSend    atomic.Int64
	BytesReceive atomic.Int64
}

type GeoSite struct {
	domain       map[string]struct{}
	fullDomain   map[string]struct{}
	regexpDomain []*regexp.Regexp
}

func NewGeoSite(data []byte) *GeoSite {
	gs := &GeoSite{
		domain:     make(map[string]struct{}),
		fullDomain: make(map[string]struct{}),
	}

	r := bufio.NewReader(bytes.NewReader(data))
	for {
		line, _, err := r.ReadLine()
		if err == io.EOF {
			break
		}

		if bytes.HasPrefix(line, []byte("full:")) {
			gs.fullDomain[string(line[5:])] = struct{}{}
			continue
		}

		if bytes.HasPrefix(line, []byte("regexp:")) {
			line = line[7:]
			re, err := regexp.Compile(string(line))
			if err != nil {
				log.Errorf("compile geosite string:%s, err:%s", string(line), err.Error())
				continue
			}
			gs.regexpDomain = append(gs.regexpDomain, re)
			continue
		}

		gs.domain[string(line)] = struct{}{}
	}

	return gs
}

func (gs *GeoSite) SiteAtCN(domain string) bool {
	domainRoot := func(_domain string) string {
		var firstDot, lastDot int
		for {
			firstDot = strings.Index(_domain, ".")
			lastDot = strings.LastIndex(_domain, ".")
			if firstDot == lastDot {
				return _domain
			}
			_domain = _domain[firstDot+1:]
		}
	}

	if _, ok := gs.fullDomain[domain]; ok {
		return true
	}

	_domain := domainRoot(domain)
	if _, ok := gs.domain[_domain]; ok {
		return true
	}

	for _, re := range gs.regexpDomain {
		if re.MatchString(domain) {
			return true
		}
	}

	return false
}

type Easyss struct {
	config         *Config
	serverIP       string
	stat           *Statistics
	localGw        string
	localDev       string
	devIndex       int
	dnsCache       *freecache.Cache
	directDNSCache *freecache.Cache
	geoipDB        *geoip2.Reader
	geosite        *GeoSite

	// the mu Mutex to protect below fields
	mu              *sync.RWMutex
	tcpPool         easypool.Pool
	socksServer     *socks5.Server
	httpProxyServer *http.Server
	closing         chan struct{}
	tun2socksStatus Tun2socksStatus
}

func New(config *Config) (*Easyss, error) {
	ss := &Easyss{
		config:         config,
		stat:           &Statistics{},
		dnsCache:       freecache.NewCache(1024 * 1024),
		directDNSCache: freecache.NewCache(1024 * 1024),
		closing:        make(chan struct{}, 1),
		mu:             &sync.RWMutex{},
	}

	db, err := geoip2.FromBytes(geoIPCNPrivate)
	if err != nil {
		panic(err)
	}
	ss.geoipDB = db
	ss.geosite = NewGeoSite(geoSiteCN)

	msg, err := ss.ServerDNSMsg()
	if err != nil {
		log.Errorf("query server dns msg err:%s", err.Error())
	}
	if msg != nil {
		ss.serverIP = msg.Answer[0].(*dns.A).A.String()
		ss.SetDNSCache(msg, true, true)
		ss.SetDNSCache(msg, true, false)
	}

	switch runtime.GOOS {
	case "linux", "windows", "darwin":
		gw, dev, err := util.SysGatewayAndDevice()
		if err != nil {
			log.Errorf("get system gateway and device err:%s", err.Error())
		}
		ss.localGw = gw
		ss.localDev = dev

		iface, err := net.InterfaceByName(dev)
		if err != nil {
			log.Errorf("interface by name err:%v", err)
			return nil, err
		}
		ss.devIndex = iface.Index
	}

	go ss.printStatistics()

	return ss, err
}

func (ss *Easyss) InitTcpPool() error {
	if ss.DisableUTLS() {
		log.Infof("uTLS is disabled")
	} else {
		log.Infof("uTLS is enabled")
	}

	factory := func() (net.Conn, error) {
		ctx, cancel := context.WithTimeout(context.Background(), ss.Timeout())
		defer cancel()

		if ss.DisableUTLS() {
			dialer := new(tls.Dialer)
			return dialer.DialContext(ctx, "tcp", ss.ServerAddr())
		}

		conn, err := net.DialTimeout("tcp", ss.ServerAddr(), ss.Timeout())
		if err != nil {
			return nil, err
		}

		uConn := utls.UClient(conn, &utls.Config{ServerName: ss.Server()}, utls.HelloChrome_Auto)
		if err := uConn.HandshakeContext(ctx); err != nil {
			return nil, err
		}

		return uConn, nil
	}

	config := &easypool.PoolConfig{
		InitialCap:  10,
		MaxCap:      50,
		MaxIdle:     10,
		Idletime:    5 * time.Minute,
		MaxLifetime: 30 * time.Minute,
		Factory:     factory,
	}
	tcpPool, err := easypool.NewHeapPool(config)
	ss.SetPool(tcpPool)

	return err
}

func (ss *Easyss) LocalPort() int {
	return ss.config.LocalPort
}

func (ss *Easyss) LocalHttpProxyPort() int {
	return ss.config.LocalPort + 1000
}

func (ss *Easyss) LocalPacPort() int {
	return ss.config.LocalPort + 1001
}

func (ss *Easyss) ServerPort() int {
	return ss.config.ServerPort
}

func (ss *Easyss) Password() string {
	return ss.config.Password
}

func (ss *Easyss) Method() string {
	return ss.config.Method
}

func (ss *Easyss) Server() string {
	return ss.config.Server
}

func (ss *Easyss) ServerIP() string {
	return ss.serverIP
}

func (ss *Easyss) ServerAddr() string {
	return fmt.Sprintf("%s:%d", ss.Server(), ss.ServerPort())
}

func (ss *Easyss) Socks5ProxyAddr() string {
	return fmt.Sprintf("socks5://%s", ss.LocalAddr())
}

func (ss *Easyss) LocalGateway() string {
	return ss.localGw
}

func (ss *Easyss) LocalDevice() string {
	return ss.localDev
}

func (ss *Easyss) LocalDeviceIndex() int {
	return ss.devIndex
}

func (ss *Easyss) Timeout() time.Duration {
	return time.Duration(ss.config.Timeout) * time.Second
}

func (ss *Easyss) LocalAddr() string {
	return fmt.Sprintf("%s:%d", "127.0.0.1", ss.LocalPort())
}

func (ss *Easyss) BindAll() bool {
	return ss.config.BindALL
}

func (ss *Easyss) DisableUTLS() bool {
	return ss.config.DisableUTLS
}

func (ss *Easyss) ConfigFilename() string {
	if ss.config.ConfigFile == "" {
		return ""
	}
	return filepath.Base(ss.config.ConfigFile)
}

func (ss *Easyss) Pool() easypool.Pool {
	ss.mu.RLock()
	defer ss.mu.RUnlock()
	return ss.tcpPool
}

func (ss *Easyss) SetSocksServer(server *socks5.Server) {
	ss.mu.Lock()
	defer ss.mu.Unlock()
	ss.socksServer = server
}

func (ss *Easyss) Tun2socksStatus() Tun2socksStatus {
	ss.mu.RLock()
	defer ss.mu.RUnlock()
	return ss.tun2socksStatus
}

func (ss *Easyss) SetTun2socksStatus(status Tun2socksStatus) {
	ss.mu.Lock()
	defer ss.mu.Unlock()
	ss.tun2socksStatus = status
}

func (ss *Easyss) Tun2socksStatusAuto() bool {
	ss.mu.RLock()
	defer ss.mu.RUnlock()
	return ss.tun2socksStatus == Tun2socksStatusAuto
}

func (ss *Easyss) Tun2socksStatusOn() bool {
	ss.mu.RLock()
	defer ss.mu.RUnlock()
	return ss.tun2socksStatus == Tun2socksStatusOn
}

func (ss *Easyss) Tun2socksStatusOff() bool {
	ss.mu.RLock()
	defer ss.mu.RUnlock()
	return ss.tun2socksStatus == Tun2socksStatusOff
}

func (ss *Easyss) SetHttpProxyServer(server *http.Server) {
	ss.mu.Lock()
	defer ss.mu.Unlock()
	ss.httpProxyServer = server
}

func (ss *Easyss) SetPool(pool easypool.Pool) {
	ss.mu.Lock()
	defer ss.mu.Unlock()
	ss.tcpPool = pool
}

func (ss *Easyss) DNSCache(name, qtype string, isDirect bool) *dns.Msg {
	var v []byte
	var err error
	if isDirect {
		v, err = ss.directDNSCache.Get([]byte(name + qtype))
	} else {
		v, err = ss.dnsCache.Get([]byte(name + qtype))
	}
	if err != nil || len(v) == 0 {
		return nil
	}

	msg := &dns.Msg{}
	if err := msg.Unpack(v); err != nil {
		return nil
	}

	return msg
}

func (ss *Easyss) RenewDNSCache(name, qtype string, isDirect bool) {
	if isDirect {
		ss.directDNSCache.Touch([]byte(name+qtype), 8*60*60)
		return
	}
	ss.dnsCache.Touch([]byte(name+qtype), 8*60*60)
}

func (ss *Easyss) SetDNSCache(msg *dns.Msg, noExpire, isDirect bool) error {
	if msg == nil {
		return nil
	}
	if len(msg.Question) == 0 {
		return nil
	}

	q := msg.Question[0]
	if q.Qtype == dns.TypeA || q.Qtype == dns.TypeAAAA {
		v, err := msg.Pack()
		if err != nil {
			return err
		}
		expireSec := 8 * 60 * 60
		if noExpire {
			expireSec = 0
		}
		key := []byte(q.Name + dns.TypeToString[q.Qtype])
		if isDirect {
			return ss.directDNSCache.Set(key, v, expireSec)
		}
		return ss.dnsCache.Set(key, v, expireSec)
	}

	return nil
}

func (ss *Easyss) ServerDNSMsg() (*dns.Msg, error) {
	c := new(dns.Client)

	m := new(dns.Msg)
	m.SetQuestion(dns.Fqdn(ss.Server()), dns.TypeA)
	m.RecursionDesired = true

	r, _, err := c.Exchange(m, "114.114.114.114:53")
	if err != nil {
		return nil, err
	}

	if r.Rcode != dns.RcodeSuccess {
		return nil, fmt.Errorf("dns query response Rcode:%v not equals RcodeSuccess", r.Rcode)
	}

	return r, nil
}

func (ss *Easyss) HostAtCN(host string) bool {
	if host == "" {
		return false
	}

	if util.IsIP(host) {
		return ss.IPAtCN(host)
	}

	return ss.geosite.SiteAtCN(host)
}

func (ss *Easyss) IPAtCN(ip string) bool {
	_ip := net.ParseIP(ip)
	if _ip == nil {
		return false
	}
	country, err := ss.geoipDB.Country(_ip)
	if err != nil {
		return false
	}

	if country.Country.IsoCode == "CN" {
		return true
	}

	return false
}

func (ss *Easyss) Close() {
	ss.mu.Lock()
	defer ss.mu.Unlock()

	if ss.tun2socksStatus != Tun2socksStatusOff {
		ss.closeTun2socks()
	}
	if ss.tcpPool != nil {
		ss.tcpPool.Close()
		ss.tcpPool = nil
	}
	if ss.httpProxyServer != nil {
		ss.httpProxyServer.Close()
		ss.httpProxyServer = nil
	}
	if ss.socksServer != nil {
		ss.socksServer.Shutdown()
		ss.socksServer = nil
	}
	if ss.closing != nil {
		close(ss.closing)
		ss.closing = nil
	}
}

func (ss *Easyss) printStatistics() {
	ss.mu.Lock()
	closing := ss.closing
	ss.mu.Unlock()
	for {
		select {
		case <-time.After(time.Hour):
			sendSize := ss.stat.BytesSend.Load() / (1024 * 1024)
			receiveSize := ss.stat.BytesReceive.Load() / (1024 * 1024)
			log.Debugf("easyss send data size: %vMB, recive data size: %vMB", sendSize, receiveSize)
		case <-closing:
			return
		}
	}
}
<|endoftext|>


<commit_before>// mounttabled is a simple mount table daemon.
package main

import (
	"flag"
	"fmt"
	"os"

	"veyron2/rt"
	"veyron2/vlog"

	"veyron/lib/signals"

	"veyron/services/mounttable/lib"
)

var (
	mountName = flag.String("name", "", "Name to mount this mountable as.  Empty means don't mount.")
	address   = flag.String("address", ":0", "Address to listen on.  Default is to use a randomly assigned port")
)

const usage = `%s is a simple mount table daemon.

Usage:

  %s [--name=<name>]

  <name>, if provided, causes the mount table to mount itself under that name.
  The name may be absolute for a remote mount table service (e.g., "/<remote mt
  address>//some/suffix") or could be relative to this process' default mount
  table (e.g., "some/suffix").
`

func Usage() {
	fmt.Fprintf(os.Stderr, usage, os.Args[0], os.Args[0])
}

func main() {
	// TODO(cnicolaou): fix Usage so that it includes the flags defined by
	// the runtime
	flag.Usage = Usage
	r := rt.Init()
	defer r.Shutdown()

	server, err := r.NewServer()
	if err != nil {
		vlog.Errorf("r.NewServer failed: %v", err)
		return
	}
	defer server.Stop()
	mtPrefix := "mt"
	if err := server.Register(mtPrefix, mounttable.NewMountTable()); err != nil {
		vlog.Errorf("server.Register failed to register mount table: %v", err)
		return
	}
	endpoint, err := server.Listen("tcp", *address)
	if err != nil {
		vlog.Errorf("server.Listen failed: %v", err)
		return
	}
	if name := *mountName; len(name) > 0 {
		if err := server.Publish(name); err != nil {
			vlog.Errorf("Publish(%v) failed: %v", name, err)
			return
		}
		vlog.Infof("Mount table service at: %v/%v (/%v/%v)", name, mtPrefix, endpoint, mtPrefix)

	} else {
		vlog.Infof("Mount table at: /%v/%v", endpoint, mtPrefix)
	}

	// Wait until signal is received.
	vlog.Info("Received signal ", <-signals.ShutdownOnSignals())
}
<commit_msg>veyron/services/mounttable/mounttabled: Allow the prefix to be specified via flag.<commit_after>// mounttabled is a simple mount table daemon.
package main

import (
	"flag"
	"fmt"
	"os"

	"veyron2/naming"
	"veyron2/rt"
	"veyron2/vlog"

	"veyron/lib/signals"

	"veyron/services/mounttable/lib"
)

var (
	mountName = flag.String("name", "", "Name to mount this mountable as.  Empty means don't mount.")
	address   = flag.String("address", ":0", "Address to listen on.  Default is to use a randomly assigned port")
	prefix    = flag.String("prefix", "mt", "The prefix to register the server at.")
)

const usage = `%s is a simple mount table daemon.

Usage:

  %s [--name=<name>]

  <name>, if provided, causes the mount table to mount itself under that name.
  The name may be absolute for a remote mount table service (e.g., "/<remote mt
  address>//some/suffix") or could be relative to this process' default mount
  table (e.g., "some/suffix").
`

func Usage() {
	fmt.Fprintf(os.Stderr, usage, os.Args[0], os.Args[0])
}

func main() {
	// TODO(cnicolaou): fix Usage so that it includes the flags defined by
	// the runtime
	flag.Usage = Usage
	r := rt.Init()
	defer r.Shutdown()

	server, err := r.NewServer()
	if err != nil {
		vlog.Errorf("r.NewServer failed: %v", err)
		return
	}
	defer server.Stop()
	if err := server.Register(*prefix, mounttable.NewMountTable()); err != nil {
		vlog.Errorf("server.Register failed to register mount table: %v", err)
		return
	}
	endpoint, err := server.Listen("tcp", *address)
	if err != nil {
		vlog.Errorf("server.Listen failed: %v", err)
		return
	}

	if name := *mountName; len(name) > 0 {
		if err := server.Publish(name); err != nil {
			vlog.Errorf("Publish(%v) failed: %v", name, err)
			return
		}
		vlog.Infof("Mount table service at: %s (%s)",
			naming.JoinAddressName(name, *prefix),
			naming.JoinAddressName(endpoint.String(), *prefix))

	} else {
		vlog.Infof("Mount table at: %s",
			naming.JoinAddressName(endpoint.String(), *prefix))
	}

	// Wait until signal is received.
	vlog.Info("Received signal ", <-signals.ShutdownOnSignals())
}
<|endoftext|>


<commit_before>package controllers

import (
	"github.com/astaxie/beego"
	"goblog/models"
	"html/template"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"time"
)

type TopicController struct {
	beego.Controller
}

// 根据分类名称查询文章
func (this *TopicController) ViewTopicByCategoryName() {

	category := this.Ctx.Input.Param(":category")

	topics := models.QueryTopicByCategoryName(category)

	categories := models.GetAllCategory() //查询所有的分类

	this.Data["Topics"] = topics
	this.Data["Categories"] = categories
	this.Data["Category"] = category

	this.TplNames = "view_topic_cat.html"
}

// 根据文章id删除文章
func (this *TopicController) DeleteTopic() {

	flash := beego.NewFlash()

	if checkAccountSession(&this.Controller) { //验证用户是否已登录
		id, err := strconv.ParseInt(this.Ctx.Input.Param(":id"), 10, 64)
		if err != nil {
			beego.Error("转换文章id失败")
			flash.Error("删除文章失败!")
			flash.Store(&this.Controller)
			return
		}
		if !models.DeleteTopic(id) {
			beego.Error("删除文章失败")
			flash.Error("删除文章失败!")
			flash.Store(&this.Controller)
			return
		}
		this.Redirect("/", 302) //删除成功回首页
		return
	} else {
		flash.Error("您尚未登录,请登录!")
		flash.Store(&this.Controller)
		this.Redirect("/login", 302) //跳转到登录页
		return
	}

}

// 修改文章页面
func (this *TopicController) ModifyTopic() {

	beego.ReadFromRequest(&this.Controller)
	flash := beego.NewFlash()
	if checkAccountSession(&this.Controller) { //验证用户是否已登录
		id, err := strconv.ParseInt(this.Ctx.Input.Param(":id"), 10, 64)
		if err != nil {
			beego.Error("获取文章id失败")
			flash.Error("获取文章id失败!")
			flash.Store(&this.Controller)
			return
		}
		topic := models.ViewTopicById(id)
		this.Data["Topic"] = topic
		this.TplNames = "modify_topic.html"
	} else {
		flash.Error("您尚未登录,请登录!")
		flash.Store(&this.Controller)
		this.Redirect("/login", 302) //跳转到登录页
		return
	}
}

// 修改文章Action
func (this *TopicController) ModifyTopicAction() {
	flash := beego.NewFlash()
	if checkAccountSession(&this.Controller) { //验证用户是否已登录

		topic := &models.Topic{}
		err := this.ParseForm(topic)

		if err != nil {
			beego.Error("收集表单数据失败!" + err.Error())
		} else {

			category := models.GetCategoryByName(topic.Category) //查询该分类名称查询该分类是否存在

			if category == nil {
				//不存在,则添加分类
				if !models.AddCategory(topic.Category) {
					beego.Error("添加文章分类失败:" + err.Error())
					flash.Error("添加文章分类失败!")
					flash.Store(&this.Controller)
					this.Redirect("/topic/modify/"+strconv.FormatInt(topic.Id, 10), 302) //重新定向到修改页面
					return
				}
			} else { //存在则修改分类
				category.CategoryName = topic.Category
				if !models.ModifyCategory(category) {
					beego.Error("修改文章分类失败:" + err.Error())
					flash.Error("修改文章分类失败!")
					flash.Store(&this.Controller)
					this.Redirect("/topic/modify/"+strconv.FormatInt(topic.Id, 10), 302) //重新定向到修改页面
					return
				}
			}

			//修改文章
			err = models.ModifyTopic(topic)
			if err != nil {
				beego.Error("修改文章失败!" + err.Error())
				flash.Error("修改文章失败!")
				flash.Store(&this.Controller)
				this.Redirect("/topic/modify/"+strconv.FormatInt(topic.Id, 10), 302) //重新定向到修改页面
				return
			}
			flash.Notice("修改文章成功!")
			flash.Store(&this.Controller)
			this.Redirect("/topic/view/"+strconv.FormatInt(topic.Id, 10), 302) //修改成功重定向到查看页面
		}
	} else {
		flash.Error("您尚未登录,请登录!")
		flash.Store(&this.Controller)
		this.Redirect("/login", 302) //跳转到登录页
		return
	}
}

// 根据文章id查看文章
func (this *TopicController) ViewTopic() {

	beego.ReadFromRequest(&this.Controller)
	flash := beego.NewFlash()

	if this.GetSession("user") != nil {
		user := this.GetSession("user").(*models.User) //从Session中获取用户信息
		this.Data["Nickname"] = user.Nickname
		this.Data["Username"] = user.Username
		this.Data["IsLogin"] = true
		this.Data["IsTopic"] = true
	}
	id, err := strconv.ParseInt(this.Ctx.Input.Param(":id"), 10, 64)
	if err != nil {
		beego.Error("获取文章id失败" + err.Error())
		flash.Error("获取文章id失败!")
		flash.Store(&this.Controller)
		return
	}
	topic := models.ViewTopicById(id)
	this.Data["Topic"] = topic
	this.TplNames = "view_topic.html"
}

// 跳转到新增页面
func (this *TopicController) Add() {

	beego.ReadFromRequest(&this.Controller)

	flash := beego.NewFlash()

	if checkAccountSession(&this.Controller) { //验证用户是否已登录
		this.Data["IsTopic"] = true
		this.Data["Title"] = "添加文章"
		this.TplNames = "add_topic.html"
		return
	} else {
		flash.Error("您尚未登录,请登录!")
		flash.Store(&this.Controller)
		this.Redirect("/login", 302) //跳转到登录页
		return
	}
}

// 添加文章内容
func (this *TopicController) AddTopic() {
	flash := beego.NewFlash()
	if checkAccountSession(&this.Controller) { //验证用户是否已登录
		topic := &models.Topic{}
		err := this.ParseForm(topic)
		if err != nil {
			beego.Error("添加文章内容失败:" + err.Error())
			flash.Error("添加文章失败!")
			flash.Store(&this.Controller)
			this.Redirect("/", 302)
			return
		}

		topic.Created = time.Now().Local() //设置创建时间
		topic.Updated = time.Now().Local() //设置更新时间

		if this.GetSession("user") != nil {
			topic.Author = this.GetSession("user").(*models.User).Nickname   //设置作者
			topic.Username = this.GetSession("user").(*models.User).Username //设置用户名
		}

		if nil == models.GetCategoryByName(topic.Category) { //查询该分类是否存在
			//不存在,则保存该分类
			if !models.AddCategory(topic.Category) {
				beego.Error("添加文章分类失败:" + err.Error())
				flash.Error("添加文章分类失败!")
				flash.Store(&this.Controller)
			}
		}

		err = models.AddTopic(topic)

		if err != nil {
			beego.Error("添加文章失败:" + err.Error())
			flash.Error("添加文章失败!")
			flash.Store(&this.Controller)
			this.Redirect("/topic/add", 302)
			return
		}

		staticpath, err := filepath.Abs("html/" + strconv.FormatInt(topic.Id, 10) + ".html")
		if err != nil {
			log.Print("获取文件的物理路径失败:" + err.Error())
		}

		file, err := os.Create(staticpath)
		if err != nil {
			log.Print("创建文件失败:" + err.Error())
		}

		tp1, err := filepath.Abs("views/header.tpl")
		tp2, err := filepath.Abs("views/view_topic.html")
		tp3, err := filepath.Abs("views/footer.tpl")
		tp4, err := filepath.Abs("views/msg.tpl")
		tp5, err := filepath.Abs("views/nav.tpl")

		if err != nil {
			log.Print("读取模板失败:" + err.Error())
		}

		var tplFuncMap template.FuncMap
		tplFuncMap = make(template.FuncMap)
		tplFuncMap["dateformat"] = beego.DateFormat
		t := template.New("view_topic.html")
		t, err = t.Funcs(tplFuncMap).ParseFiles(tp1, tp2, tp3, tp4, tp5)

		if err != nil {
			log.Print("解析模板失败:" + err.Error())
		}

		data := map[string]interface{}{
			"Title": "title",
			"Topic": topic,
		}

		err = t.Execute(file, data)
		if err != nil {
			log.Print("解析模板失败:" + err.Error())
		}

		this.Redirect("/", 302) //添加成功到文章列表页
		return
	} else {
		flash.Error("您尚未登录,请登录!")
		flash.Store(&this.Controller)
		this.Redirect("/login", 302) //跳转到登录页
		return
	}

}
<commit_msg>简化模板静态化代码<commit_after>package controllers

import (
	"github.com/astaxie/beego"
	"goblog/models"
	"html/template"
	"log"
	"os"
	"path/filepath"
	"strconv"
	"time"
)

type TopicController struct {
	beego.Controller
}

// 根据分类名称查询文章
func (this *TopicController) ViewTopicByCategoryName() {

	category := this.Ctx.Input.Param(":category")

	topics := models.QueryTopicByCategoryName(category)

	categories := models.GetAllCategory() //查询所有的分类

	this.Data["Topics"] = topics
	this.Data["Categories"] = categories
	this.Data["Category"] = category

	this.TplNames = "view_topic_cat.html"
}

// 根据文章id删除文章
func (this *TopicController) DeleteTopic() {

	flash := beego.NewFlash()

	if checkAccountSession(&this.Controller) { //验证用户是否已登录
		id, err := strconv.ParseInt(this.Ctx.Input.Param(":id"), 10, 64)
		if err != nil {
			beego.Error("转换文章id失败")
			flash.Error("删除文章失败!")
			flash.Store(&this.Controller)
			return
		}
		if !models.DeleteTopic(id) {
			beego.Error("删除文章失败")
			flash.Error("删除文章失败!")
			flash.Store(&this.Controller)
			return
		}
		this.Redirect("/", 302) //删除成功回首页
		return
	} else {
		flash.Error("您尚未登录,请登录!")
		flash.Store(&this.Controller)
		this.Redirect("/login", 302) //跳转到登录页
		return
	}

}

// 修改文章页面
func (this *TopicController) ModifyTopic() {

	beego.ReadFromRequest(&this.Controller)
	flash := beego.NewFlash()
	if checkAccountSession(&this.Controller) { //验证用户是否已登录
		id, err := strconv.ParseInt(this.Ctx.Input.Param(":id"), 10, 64)
		if err != nil {
			beego.Error("获取文章id失败")
			flash.Error("获取文章id失败!")
			flash.Store(&this.Controller)
			return
		}
		topic := models.ViewTopicById(id)
		this.Data["Topic"] = topic
		this.TplNames = "modify_topic.html"
	} else {
		flash.Error("您尚未登录,请登录!")
		flash.Store(&this.Controller)
		this.Redirect("/login", 302) //跳转到登录页
		return
	}
}

// 修改文章Action
func (this *TopicController) ModifyTopicAction() {
	flash := beego.NewFlash()
	if checkAccountSession(&this.Controller) { //验证用户是否已登录

		topic := &models.Topic{}
		err := this.ParseForm(topic)

		if err != nil {
			beego.Error("收集表单数据失败!" + err.Error())
		} else {

			category := models.GetCategoryByName(topic.Category) //查询该分类名称查询该分类是否存在

			if category == nil {
				//不存在,则添加分类
				if !models.AddCategory(topic.Category) {
					beego.Error("添加文章分类失败:" + err.Error())
					flash.Error("添加文章分类失败!")
					flash.Store(&this.Controller)
					this.Redirect("/topic/modify/"+strconv.FormatInt(topic.Id, 10), 302) //重新定向到修改页面
					return
				}
			} else { //存在则修改分类
				category.CategoryName = topic.Category
				if !models.ModifyCategory(category) {
					beego.Error("修改文章分类失败:" + err.Error())
					flash.Error("修改文章分类失败!")
					flash.Store(&this.Controller)
					this.Redirect("/topic/modify/"+strconv.FormatInt(topic.Id, 10), 302) //重新定向到修改页面
					return
				}
			}

			//修改文章
			err = models.ModifyTopic(topic)
			if err != nil {
				beego.Error("修改文章失败!" + err.Error())
				flash.Error("修改文章失败!")
				flash.Store(&this.Controller)
				this.Redirect("/topic/modify/"+strconv.FormatInt(topic.Id, 10), 302) //重新定向到修改页面
				return
			}
			flash.Notice("修改文章成功!")
			flash.Store(&this.Controller)
			this.Redirect("/topic/view/"+strconv.FormatInt(topic.Id, 10), 302) //修改成功重定向到查看页面
		}
	} else {
		flash.Error("您尚未登录,请登录!")
		flash.Store(&this.Controller)
		this.Redirect("/login", 302) //跳转到登录页
		return
	}
}

// 根据文章id查看文章
func (this *TopicController) ViewTopic() {

	beego.ReadFromRequest(&this.Controller)
	flash := beego.NewFlash()

	if this.GetSession("user") != nil {
		user := this.GetSession("user").(*models.User) //从Session中获取用户信息
		this.Data["Nickname"] = user.Nickname
		this.Data["Username"] = user.Username
		this.Data["IsLogin"] = true
		this.Data["IsTopic"] = true
	}
	id, err := strconv.ParseInt(this.Ctx.Input.Param(":id"), 10, 64)
	if err != nil {
		beego.Error("获取文章id失败" + err.Error())
		flash.Error("获取文章id失败!")
		flash.Store(&this.Controller)
		return
	}
	topic := models.ViewTopicById(id)
	this.Data["Topic"] = topic
	this.TplNames = "view_topic.html"
}

// 跳转到新增页面
func (this *TopicController) Add() {

	beego.ReadFromRequest(&this.Controller)

	flash := beego.NewFlash()

	if checkAccountSession(&this.Controller) { //验证用户是否已登录
		this.Data["IsTopic"] = true
		this.Data["Title"] = "添加文章"
		this.TplNames = "add_topic.html"
		return
	} else {
		flash.Error("您尚未登录,请登录!")
		flash.Store(&this.Controller)
		this.Redirect("/login", 302) //跳转到登录页
		return
	}
}

// 添加文章内容
func (this *TopicController) AddTopic() {
	flash := beego.NewFlash()
	if checkAccountSession(&this.Controller) { //验证用户是否已登录
		topic := &models.Topic{}
		err := this.ParseForm(topic)
		if err != nil {
			beego.Error("添加文章内容失败:" + err.Error())
			flash.Error("添加文章失败!")
			flash.Store(&this.Controller)
			this.Redirect("/", 302)
			return
		}

		topic.Created = time.Now().Local() //设置创建时间
		topic.Updated = time.Now().Local() //设置更新时间

		if this.GetSession("user") != nil {
			topic.Author = this.GetSession("user").(*models.User).Nickname   //设置作者
			topic.Username = this.GetSession("user").(*models.User).Username //设置用户名
		}

		if nil == models.GetCategoryByName(topic.Category) { //查询该分类是否存在
			//不存在,则保存该分类
			if !models.AddCategory(topic.Category) {
				beego.Error("添加文章分类失败:" + err.Error())
				flash.Error("添加文章分类失败!")
				flash.Store(&this.Controller)
			}
		}

		err = models.AddTopic(topic)

		if err != nil {
			beego.Error("添加文章失败:" + err.Error())
			flash.Error("添加文章失败!")
			flash.Store(&this.Controller)
			this.Redirect("/topic/add", 302)
			return
		}

		staticpath, err := filepath.Abs("html/" + strconv.FormatInt(topic.Id, 10) + ".html")
		if err != nil {
			log.Print("获取文件的物理路径失败:" + err.Error())
		}

		file, err := os.Create(staticpath)
		if err != nil {
			log.Print("创建文件失败:" + err.Error())
		}

		// tp1, err := filepath.Abs("views/header.tpl")
		// tp2, err := filepath.Abs("views/view_topic.html")
		// tp3, err := filepath.Abs("views/footer.tpl")
		// tp4, err := filepath.Abs("views/msg.tpl")
		// tp5, err := filepath.Abs("views/nav.tpl")

		if err != nil {
			log.Print("读取模板失败:" + err.Error())
		}

		tplFuncMap := make(template.FuncMap)
		tplFuncMap["dateformat"] = beego.DateFormat //注册模板中使用到的模板函数dateformat

		t := template.New("view_topic.html") //此处的view_topic.html为具体的模板名称
		t = t.Funcs(tplFuncMap)              /*.ParseFiles(tp1, tp2, tp3, tp4, tp5)*/
		t, err = t.ParseGlob("views/*")      //模板存放路径 将会匹配到views/目录下的所有文件

		if err != nil {
			log.Print("解析模板失败:" + err.Error())
		}

		data := map[string]interface{}{
			"Title": "title",
			"Topic": topic,
			"Time":  time.Now().Local(),
		}

		err = t.Execute(file, data)
		if err != nil {
			log.Print("解析模板失败:" + err.Error())
		}

		this.Redirect("/", 302) //添加成功到文章列表页
		return
	} else {
		flash.Error("您尚未登录,请登录!")
		flash.Store(&this.Controller)
		this.Redirect("/login", 302) //跳转到登录页
		return
	}

}
<|endoftext|>