<commit_before>package core

/*
#cgo CFLAGS: -I./src/include
#include "lwip/tcp.h"
*/
import "C"
import (
	"context"
	"errors"
	"fmt"
	// "log"
	"math/rand"
	"net"
	"sync"
	"unsafe"
)

type tcpConn struct {
	sync.Mutex

	pcb         *C.struct_tcp_pcb
	handler     ConnectionHandler
	network     string
	remoteAddr  net.Addr
	localAddr   net.Addr
	connKeyArg  unsafe.Pointer
	connKey     uint32
	closing     bool
	localClosed bool
	aborting    bool
	ctx         context.Context
	cancel      context.CancelFunc

	// Data from remote not yet write to local will buffer into this channel.
	localWriteCh    chan []byte
	localWriteSubCh chan []byte
}

// func checkTCPConns() {
// 	tcpConns.Range(func(_, c interface{}) bool {
// 		state := c.(*tcpConn).pcb.state
// 		if c.(*tcpConn).pcb == nil ||
// 			state == C.CLOSED ||
// 			state == C.CLOSE_WAIT {
// 			c.(*tcpConn).Release()
// 		}
// 		return true
// 	})
// }

func NewTCPConnection(pcb *C.struct_tcp_pcb, handler ConnectionHandler) (Connection, error) {
	// prepare key
	connKeyArg := NewConnKeyArg()
	connKey := rand.Uint32()
	SetConnKeyVal(unsafe.Pointer(connKeyArg), connKey)

	if tcpConnectionHandler == nil {
		return nil, errors.New("no registered TCP connection handlers found")
	}

	ctx, cancel := context.WithCancel(context.Background())

	conn := &tcpConn{
		pcb:             pcb,
		handler:         handler,
		network:         "tcp",
		localAddr:       ParseTCPAddr(IPAddrNTOA(pcb.remote_ip), uint16(pcb.remote_port)),
		remoteAddr:      ParseTCPAddr(IPAddrNTOA(pcb.local_ip), uint16(pcb.local_port)),
		connKeyArg:      connKeyArg,
		connKey:         connKey,
		closing:         false,
		localClosed:     false,
		aborting:        false,
		ctx:             ctx,
		cancel:          cancel,
		localWriteCh:    make(chan []byte, 32),
		localWriteSubCh: make(chan []byte, 1),
	}

	// Associate conn with key and save to the global map.
	tcpConns.Store(connKey, conn)

	// go checkTCPConns()

	// Pass the key as arg for subsequent tcp callbacks.
	C.tcp_arg(pcb, unsafe.Pointer(connKeyArg))

	SetTCPRecvCallback(pcb)
	SetTCPSentCallback(pcb)
	SetTCPErrCallback(pcb)
	SetTCPPollCallback(pcb, C.u8_t(1)) // interval 1 means Poll will be called twice a second

	err := handler.Connect(conn, conn.RemoteAddr())
	if err != nil {
		return nil, err
	}
	return conn, nil
}

func (conn *tcpConn) RemoteAddr() net.Addr {
	return conn.remoteAddr
}

func (conn *tcpConn) LocalAddr() net.Addr {
	return conn.localAddr
}

func (conn *tcpConn) Receive(data []byte) error {
	if conn.isClosing() {
		return errors.New(fmt.Sprintf("connection %v->%v was closed by remote", conn.LocalAddr(), conn.RemoteAddr()))
	}

	err := conn.handler.DidReceive(conn, data)
	if err != nil {
		return errors.New(fmt.Sprintf("write proxy failed: %v", err))
	}

	C.tcp_recved(conn.pcb, C.u16_t(len(data)))

	return nil
}

func (conn *tcpConn) tryWriteLocal() {
	lwipMutex.Lock()
	defer lwipMutex.Unlock()

Loop:
	for {
		// Using 2 select to ensure data in localWriteSubCh will be drained first.
		select {
		case data := <-conn.localWriteSubCh:
			written, err := conn.tcpWrite(data)
			if !written || err != nil {
				// Data not written, buffer again.
				conn.localWriteSubCh <- data
				break Loop
			}
		default:
		}

		select {
		case data := <-conn.localWriteSubCh:
			written, err := conn.tcpWrite(data)
			if !written || err != nil {
				// Data not written, buffer again.
				conn.localWriteSubCh <- data
				break Loop
			}
		case data := <-conn.localWriteCh:
			written, err := conn.tcpWrite(data)
			if !written || err != nil {
				// If writing is not success, buffer to the sub channel, and next time
				// we try to read from the sub channel first. Using a sub channel here
				// because the data must be sent in correct order and we have no way
				// to prepend data to the head of a channel.
				conn.localWriteSubCh <- data
				break Loop
			}
		default:
			break Loop
		}
	}

	// Actually send data.
	C.tcp_output(conn.pcb)
	// err := C.tcp_output(conn.pcb)
	// if err != C.ERR_OK {
	// 	log.Printf("tcp_output error with lwip error code: %v", int(err))
	// }
}

// tcpWrite enqueues data to snd_buf, and treats ERR_MEM returned by tcp_write not an error,
// but instead tells the caller that data is not successfully enqueued, and should try
// again another time. By calling this function, the lwIP thread is assumed to be already
// locked by the caller.
func (conn *tcpConn) tcpWrite(data []byte) (bool, error) {
	if len(data) <= int(conn.pcb.snd_buf) {
		// Enqueue data, data copy here! Copying is required because lwIP must keep the data until they
		// are acknowledged (receiving ACK segments) by other hosts for retransmission purposes, it's
		// not obvious how to implement zero-copy here.
		err := C.tcp_write(conn.pcb, unsafe.Pointer(&data[0]), C.u16_t(len(data)), C.TCP_WRITE_FLAG_COPY)
		if err == C.ERR_OK {
			return true, nil
		} else if err != C.ERR_MEM {
			return false, errors.New(fmt.Sprintf("lwip tcp_write failed with error code: %v", int(err)))
		}
	}
	return false, nil
}

func (conn *tcpConn) Write(data []byte) (int, error) {
	if conn.isLocalClosed() {
		return 0, errors.New(fmt.Sprintf("connection %v->%v was closed by local", conn.LocalAddr(), conn.RemoteAddr()))
	}

	var written = false
	var err error

	// If there isn't any pending data left, we can try to write the data first to avoid one copy,
	// if there is pending data not yet sent, we must copy and buffer the data in order to maintain
	// the transmission order.
	if !conn.hasPendingLocalData() {
		lwipMutex.Lock()
		written, err = conn.tcpWrite(data)
		lwipMutex.Unlock()
		if err != nil {
			return 0, err
		}
	}

	if !written {
		select {
		// Buffer the data here and try sending it later, one could set a smaller localWriteCh size
		// to limit data copying times and memory usage, by sacrificing performance. But writing data
		// to local is quite fast, thus it should be safe even has a size of 1 localWriteCh.
		case conn.localWriteCh <- append([]byte(nil), data...): // data copy here!
		case <-conn.ctx.Done():
			return 0, conn.ctx.Err()
		}
	}

	// Try to send pending data if any, and call tcp_output().
	go conn.tryWriteLocal()

	return len(data), nil
}

func (conn *tcpConn) Sent(len uint16) error {
	conn.handler.DidSend(conn, len)
	// Some packets are acknowledged by local client, check if any pending data to send.
	return conn.CheckState()
}

func (conn *tcpConn) isClosing() bool {
	conn.Lock()
	defer conn.Unlock()
	return conn.closing
}

func (conn *tcpConn) isAborting() bool {
	conn.Lock()
	defer conn.Unlock()
	return conn.aborting
}

func (conn *tcpConn) isLocalClosed() bool {
	conn.Lock()
	defer conn.Unlock()
	return conn.localClosed
}

func (conn *tcpConn) hasPendingLocalData() bool {
	if len(conn.localWriteCh) > 0 || len(conn.localWriteSubCh) > 0 {
		return true
	}
	return false
}

func (conn *tcpConn) CheckState() error {
	// Still have data to send
	if conn.hasPendingLocalData() && !conn.isLocalClosed() {
		go conn.tryWriteLocal()
		// Return and wait for the Sent() callback to be called, and then check again.
		return NewLWIPError(LWIP_ERR_OK)
	}

	if conn.isClosing() || conn.isLocalClosed() {
		conn.closeInternal()
	}

	if conn.isAborting() {
		conn.abortInternal()
		return NewLWIPError(LWIP_ERR_ABRT)
	}

	return NewLWIPError(LWIP_ERR_OK)
}

func (conn *tcpConn) Close() error {
	conn.Lock()
	defer conn.Unlock()

	// Close maybe called outside of lwIP thread, we should not call tcp_close() in this
	// function, instead just make a flag to indicate we are closing the connection.
	conn.closing = true
	return nil
}

func (conn *tcpConn) setLocalClosed() error {
	conn.Lock()
	defer conn.Unlock()

	conn.localClosed = true
	return nil
}

func (conn *tcpConn) closeInternal() error {
	C.tcp_arg(conn.pcb, nil)
	C.tcp_recv(conn.pcb, nil)
	C.tcp_sent(conn.pcb, nil)
	C.tcp_err(conn.pcb, nil)
	C.tcp_poll(conn.pcb, nil, 0)

	conn.Release()

	conn.cancel()

	err := C.tcp_close(conn.pcb)
	if err == C.ERR_OK {
		return nil
	} else {
		return errors.New(fmt.Sprint("close TCP connection failed, lwip error code %d", int(err)))
	}
}

func (conn *tcpConn) abortInternal() {
	// log.Printf("abort TCP connection %v->%v", conn.LocalAddr(), conn.RemoteAddr())
	conn.Release()
	C.tcp_abort(conn.pcb)
}

func (conn *tcpConn) Abort() {
	conn.Lock()
	defer conn.Unlock()

	conn.aborting = true
}

// The corresponding pcb is already freed when this callback is called
func (conn *tcpConn) Err(err error) {
	// log.Printf("error on TCP connection %v->%v: %v", conn.LocalAddr(), conn.RemoteAddr(), err)
	conn.Release()
	conn.cancel()
	conn.handler.DidClose(conn)
}

func (conn *tcpConn) LocalDidClose() error {
	// log.Printf("local close TCP connection %v->%v", conn.LocalAddr(), conn.RemoteAddr())
	conn.handler.LocalDidClose(conn)
	conn.setLocalClosed()    // flag closing
	return conn.CheckState() // check pending data
}

func (conn *tcpConn) Release() {
	if _, found := tcpConns.Load(conn.connKey); found {
		FreeConnKeyArg(conn.connKeyArg)
		tcpConns.Delete(conn.connKey)
	}
	// log.Printf("ended TCP connection %v->%v", conn.LocalAddr(), conn.RemoteAddr())
}

func (conn *tcpConn) Poll() error {
	return conn.CheckState()
}
<commit_msg>comment<commit_after>package core

/*
#cgo CFLAGS: -I./src/include
#include "lwip/tcp.h"
*/
import "C"
import (
	"context"
	"errors"
	"fmt"
	// "log"
	"math/rand"
	"net"
	"sync"
	"unsafe"
)

type tcpConn struct {
	sync.Mutex

	pcb         *C.struct_tcp_pcb
	handler     ConnectionHandler
	network     string
	remoteAddr  net.Addr
	localAddr   net.Addr
	connKeyArg  unsafe.Pointer
	connKey     uint32
	closing     bool
	localClosed bool
	aborting    bool
	ctx         context.Context
	cancel      context.CancelFunc

	// Data from remote not yet write to local will buffer into this channel.
	localWriteCh    chan []byte
	localWriteSubCh chan []byte
}

// func checkTCPConns() {
// 	tcpConns.Range(func(_, c interface{}) bool {
// 		state := c.(*tcpConn).pcb.state
// 		if c.(*tcpConn).pcb == nil ||
// 			state == C.CLOSED ||
// 			state == C.CLOSE_WAIT {
// 			c.(*tcpConn).Release()
// 		}
// 		return true
// 	})
// }

func NewTCPConnection(pcb *C.struct_tcp_pcb, handler ConnectionHandler) (Connection, error) {
	// prepare key
	connKeyArg := NewConnKeyArg()
	connKey := rand.Uint32()
	SetConnKeyVal(unsafe.Pointer(connKeyArg), connKey)

	if tcpConnectionHandler == nil {
		return nil, errors.New("no registered TCP connection handlers found")
	}

	ctx, cancel := context.WithCancel(context.Background())

	conn := &tcpConn{
		pcb:             pcb,
		handler:         handler,
		network:         "tcp",
		localAddr:       ParseTCPAddr(IPAddrNTOA(pcb.remote_ip), uint16(pcb.remote_port)),
		remoteAddr:      ParseTCPAddr(IPAddrNTOA(pcb.local_ip), uint16(pcb.local_port)),
		connKeyArg:      connKeyArg,
		connKey:         connKey,
		closing:         false,
		localClosed:     false,
		aborting:        false,
		ctx:             ctx,
		cancel:          cancel,
		localWriteCh:    make(chan []byte, 32),
		localWriteSubCh: make(chan []byte, 1),
	}

	// Associate conn with key and save to the global map.
	tcpConns.Store(connKey, conn)

	// go checkTCPConns()

	// Pass the key as arg for subsequent tcp callbacks.
	C.tcp_arg(pcb, unsafe.Pointer(connKeyArg))

	SetTCPRecvCallback(pcb)
	SetTCPSentCallback(pcb)
	SetTCPErrCallback(pcb)
	SetTCPPollCallback(pcb, C.u8_t(1)) // interval 1 means Poll will be called twice a second

	err := handler.Connect(conn, conn.RemoteAddr())
	if err != nil {
		return nil, err
	}
	return conn, nil
}

func (conn *tcpConn) RemoteAddr() net.Addr {
	return conn.remoteAddr
}

func (conn *tcpConn) LocalAddr() net.Addr {
	return conn.localAddr
}

func (conn *tcpConn) Receive(data []byte) error {
	if conn.isClosing() {
		return errors.New(fmt.Sprintf("connection %v->%v was closed by remote", conn.LocalAddr(), conn.RemoteAddr()))
	}

	err := conn.handler.DidReceive(conn, data)
	if err != nil {
		return errors.New(fmt.Sprintf("write proxy failed: %v", err))
	}

	C.tcp_recved(conn.pcb, C.u16_t(len(data)))

	return nil
}

func (conn *tcpConn) tryWriteLocal() {
	lwipMutex.Lock()
	defer lwipMutex.Unlock()

Loop:
	for {
		// Using 2 select to ensure data in localWriteSubCh will be drained first.
		select {
		case data := <-conn.localWriteSubCh:
			written, err := conn.tcpWrite(data)
			if !written || err != nil {
				// Data not written, buffer again.
				conn.localWriteSubCh <- data
				break Loop
			}
		default:
		}

		select {
		case data := <-conn.localWriteSubCh:
			written, err := conn.tcpWrite(data)
			if !written || err != nil {
				// Data not written, buffer again.
				conn.localWriteSubCh <- data
				break Loop
			}
		case data := <-conn.localWriteCh:
			written, err := conn.tcpWrite(data)
			if !written || err != nil {
				// If writing is not success, buffer to the sub channel, and next time
				// we try to read from the sub channel first. Using a sub channel here
				// because the data must be sent in correct order and we have no way
				// to prepend data to the head of a channel.
				conn.localWriteSubCh <- data
				break Loop
			}
		default:
			break Loop
		}
	}

	// Actually send data.
	C.tcp_output(conn.pcb)
	// err := C.tcp_output(conn.pcb)
	// if err != C.ERR_OK {
	// 	log.Printf("tcp_output error with lwip error code: %v", int(err))
	// }
}

// tcpWrite enqueues data to snd_buf, and treats ERR_MEM returned by tcp_write not an error,
// but instead tells the caller that data is not successfully enqueued, and should try
// again another time. By calling this function, the lwIP thread is assumed to be already
// locked by the caller.
func (conn *tcpConn) tcpWrite(data []byte) (bool, error) {
	if len(data) <= int(conn.pcb.snd_buf) {
		// Enqueue data, data copy here! Copying is required because lwIP must keep the data until they
		// are acknowledged (receiving ACK segments) by other hosts for retransmission purposes, it's
		// not obvious how to implement zero-copy here.
		err := C.tcp_write(conn.pcb, unsafe.Pointer(&data[0]), C.u16_t(len(data)), C.TCP_WRITE_FLAG_COPY)
		if err == C.ERR_OK {
			return true, nil
		} else if err != C.ERR_MEM {
			return false, errors.New(fmt.Sprintf("lwip tcp_write failed with error code: %v", int(err)))
		}
	}
	return false, nil
}

func (conn *tcpConn) Write(data []byte) (int, error) {
	if conn.isLocalClosed() {
		return 0, errors.New(fmt.Sprintf("connection %v->%v was closed by local", conn.LocalAddr(), conn.RemoteAddr()))
	}

	var written = false
	var err error

	// If there isn't any pending data left, we can try to write the data first to avoid one copy,
	// if there is pending data not yet sent, we must copy and buffer the data in order to maintain
	// the transmission order.
	if !conn.hasPendingLocalData() {
		lwipMutex.Lock()
		written, err = conn.tcpWrite(data)
		lwipMutex.Unlock()
		if err != nil {
			return 0, err
		}
	}

	if !written {
		select {
		// Buffer the data here and try sending it later, one could set a smaller localWriteCh size
		// to limit data copying times and memory usage, by sacrificing performance. But writing data
		// to local is quite fast, thus it should be safe even has a size of 1 localWriteCh.
		case conn.localWriteCh <- append([]byte(nil), data...): // data copy here!
		case <-conn.ctx.Done():
			return 0, conn.ctx.Err()
		}
	}

	// Try to send pending data if any, and call tcp_output().
	go conn.tryWriteLocal()

	return len(data), nil
}

func (conn *tcpConn) Sent(len uint16) error {
	conn.handler.DidSend(conn, len)
	// Some packets are acknowledged by local client, check if any pending data to send.
	return conn.CheckState()
}

func (conn *tcpConn) isClosing() bool {
	conn.Lock()
	defer conn.Unlock()
	return conn.closing
}

func (conn *tcpConn) isAborting() bool {
	conn.Lock()
	defer conn.Unlock()
	return conn.aborting
}

func (conn *tcpConn) isLocalClosed() bool {
	conn.Lock()
	defer conn.Unlock()
	return conn.localClosed
}

func (conn *tcpConn) hasPendingLocalData() bool {
	if len(conn.localWriteCh) > 0 || len(conn.localWriteSubCh) > 0 {
		return true
	}
	return false
}

func (conn *tcpConn) CheckState() error {
	// Still have data to send
	if conn.hasPendingLocalData() && !conn.isLocalClosed() {
		go conn.tryWriteLocal()
		// Return and wait for the Sent() callback to be called, and then check again.
		return NewLWIPError(LWIP_ERR_OK)
	}

	if conn.isClosing() || conn.isLocalClosed() {
		conn.closeInternal()
	}

	if conn.isAborting() {
		conn.abortInternal()
		return NewLWIPError(LWIP_ERR_ABRT)
	}

	return NewLWIPError(LWIP_ERR_OK)
}

func (conn *tcpConn) Close() error {
	conn.Lock()
	defer conn.Unlock()

	// Close maybe called outside of lwIP thread, we should not call tcp_close() in this
	// function, instead just make a flag to indicate we are closing the connection.
	conn.closing = true
	return nil
}

func (conn *tcpConn) setLocalClosed() error {
	conn.Lock()
	defer conn.Unlock()

	conn.localClosed = true
	return nil
}

func (conn *tcpConn) closeInternal() error {
	C.tcp_arg(conn.pcb, nil)
	C.tcp_recv(conn.pcb, nil)
	C.tcp_sent(conn.pcb, nil)
	C.tcp_err(conn.pcb, nil)
	C.tcp_poll(conn.pcb, nil, 0)

	conn.Release()

	conn.cancel()

	// TODO: may return ERR_MEM if no memory to allocate segments use for closing the conn,
	// should check and try again in Sent() for Poll() callbacks.
	err := C.tcp_close(conn.pcb)
	if err == C.ERR_OK {
		return nil
	} else {
		return errors.New(fmt.Sprint("close TCP connection failed, lwip error code %d", int(err)))
	}
}

func (conn *tcpConn) abortInternal() {
	// log.Printf("abort TCP connection %v->%v", conn.LocalAddr(), conn.RemoteAddr())
	conn.Release()
	C.tcp_abort(conn.pcb)
}

func (conn *tcpConn) Abort() {
	conn.Lock()
	defer conn.Unlock()

	conn.aborting = true
}

// The corresponding pcb is already freed when this callback is called
func (conn *tcpConn) Err(err error) {
	// log.Printf("error on TCP connection %v->%v: %v", conn.LocalAddr(), conn.RemoteAddr(), err)
	conn.Release()
	conn.cancel()
	conn.handler.DidClose(conn)
}

func (conn *tcpConn) LocalDidClose() error {
	// log.Printf("local close TCP connection %v->%v", conn.LocalAddr(), conn.RemoteAddr())
	conn.handler.LocalDidClose(conn)
	conn.setLocalClosed()    // flag closing
	return conn.CheckState() // check pending data
}

func (conn *tcpConn) Release() {
	if _, found := tcpConns.Load(conn.connKey); found {
		FreeConnKeyArg(conn.connKeyArg)
		tcpConns.Delete(conn.connKey)
	}
	// log.Printf("ended TCP connection %v->%v", conn.LocalAddr(), conn.RemoteAddr())
}

func (conn *tcpConn) Poll() error {
	return conn.CheckState()
}
<|endoftext|>


<commit_before>// SPDX-License-Identifier: Apache-2.0
// Copyright Authors of Cilium

package endpoint

import (
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"sync/atomic"
	"unsafe"

	"github.com/sirupsen/logrus"
	"gopkg.in/natefinch/lumberjack.v2"

	"github.com/cilium/cilium/pkg/logging"
	"github.com/cilium/cilium/pkg/logging/logfields"
	"github.com/cilium/cilium/pkg/option"
)

var (
	log       = logging.DefaultLogger.WithField(logfields.LogSubsys, subsystem)
	policyLog = logrus.New()

	policyLogOnce sync.Once
)

const (
	subsystem = "endpoint"

	fieldRegenLevel = "regeneration-level"
)

// getLogger returns a logrus object with EndpointID, containerID and the Endpoint
// revision fields.
func (e *Endpoint) getLogger() *logrus.Entry {
	v := atomic.LoadPointer(&e.logger)
	return (*logrus.Entry)(v)
}

// getPolicyLogger returns a logger to be used for policy update debugging, or nil,
// if not configured.
func (e *Endpoint) getPolicyLogger() *logrus.Entry {
	v := atomic.LoadPointer(&e.policyLogger)
	return (*logrus.Entry)(v)
}

// PolicyDebug logs the 'msg' with 'fields' if policy debug logging is enabled.
func (e *Endpoint) PolicyDebug(fields logrus.Fields, msg string) {
	if dbgLog := e.getPolicyLogger(); dbgLog != nil {
		dbgLog.WithFields(fields).Debug(msg)
	}
}

// Logger returns a logrus object with EndpointID, containerID and the Endpoint
// revision fields. The caller must specify their subsystem.
func (e *Endpoint) Logger(subsystem string) *logrus.Entry {
	if e == nil {
		return log.WithField(logfields.LogSubsys, subsystem)
	}

	return e.getLogger().WithField(logfields.LogSubsys, subsystem)
}

// UpdateLogger creates a logger instance specific to this endpoint. It will
// create a custom Debug logger for this endpoint when the option on it is set.
// If fields is not nil only the those specific fields will be updated in the
// endpoint's logger, otherwise a full update of those fields is executed.
//
// Note: You must hold Endpoint.mutex.Lock() to synchronize logger pointer
// updates if the endpoint is already exposed. Callers that create new
// endopoints do not need locks to call this.
func (e *Endpoint) UpdateLogger(fields map[string]interface{}) {
	e.updatePolicyLogger(fields)
	v := atomic.LoadPointer(&e.logger)
	epLogger := (*logrus.Entry)(v)
	if fields != nil && epLogger != nil {
		newLogger := epLogger.WithFields(fields)
		atomic.StorePointer(&e.logger, unsafe.Pointer(newLogger))
		return
	}

	// We need to update if
	// - e.logger is nil (this happens on the first ever call to UpdateLogger via
	//   Logger above). This clause has to come first to guard the others.
	// - If any of EndpointID, containerID or policyRevision are different on the
	//   endpoint from the logger.
	// - The debug option on the endpoint is true, and the logger is not debug,
	//   or vice versa.
	shouldUpdate := epLogger == nil || (e.Options != nil &&
		e.Options.IsEnabled(option.Debug) != (epLogger.Level == logrus.DebugLevel))

	// do nothing if we do not need an update
	if !shouldUpdate {
		return
	}

	// default to using the log var set above
	baseLogger := log.Logger

	// If this endpoint is set to debug ensure it will print debug by giving it
	// an independent logger
	if e.Options != nil && e.Options.IsEnabled(option.Debug) {
		baseLogger = logging.InitializeDefaultLogger()
		baseLogger.SetLevel(logrus.DebugLevel)
	} else {
		// Debug mode takes priority; if not in debug, check what log level user
		// has set and set the endpoint's log to log at that level.
		baseLogger.SetLevel(logging.DefaultLogger.Level)
	}

	// When adding new fields, make sure they are abstracted by a setter
	// and update the logger when the value is set.
	l := baseLogger.WithFields(logrus.Fields{
		logfields.LogSubsys:              subsystem,
		logfields.EndpointID:             e.ID,
		logfields.ContainerID:            e.getShortContainerID(),
		logfields.DatapathPolicyRevision: e.policyRevision,
		logfields.DesiredPolicyRevision:  e.nextPolicyRevision,
		logfields.IPv4:                   e.IPv4.String(),
		logfields.IPv6:                   e.IPv6.String(),
		logfields.K8sPodName:             e.getK8sNamespaceAndPodName(),
	})

	if e.SecurityIdentity != nil {
		l = l.WithField(logfields.Identity, e.SecurityIdentity.ID.StringID())
	}

	atomic.StorePointer(&e.logger, unsafe.Pointer(l))
}

// Only to be called from UpdateLogger() above
func (e *Endpoint) updatePolicyLogger(fields map[string]interface{}) {
	pv := atomic.LoadPointer(&e.policyLogger)
	policyLogger := (*logrus.Entry)(pv)
	// e.Options check needed for unit testing.
	if policyLogger == nil && e.Options != nil && e.Options.IsEnabled(option.DebugPolicy) {
		policyLogOnce.Do(func() {
			maxSize := 10 // 10 MB
			if ms := os.Getenv("CILIUM_DBG_POLICY_LOG_MAX_SIZE"); ms != "" {
				if ms, err := strconv.Atoi(ms); err == nil {
					maxSize = ms
				}
			}
			maxBackups := 3
			if mb := os.Getenv("CILIUM_DBG_POLICY_LOG_MAX_BACKUPS"); mb != "" {
				if mb, err := strconv.Atoi(mb); err == nil {
					maxBackups = mb
				}
			}
			lumberjackLogger := &lumberjack.Logger{
				Filename:   filepath.Join(option.Config.StateDir, "endpoint-policy.log"),
				MaxSize:    maxSize,
				MaxBackups: maxBackups,
				MaxAge:     28, // days
				LocalTime:  true,
				Compress:   true,
			}
			policyLog.SetOutput(lumberjackLogger)
			policyLog.SetLevel(logrus.DebugLevel)
		})
		policyLogger = logrus.NewEntry(policyLog)
	}
	if policyLogger == nil || e.Options == nil {
		return
	}

	if !e.Options.IsEnabled(option.DebugPolicy) {
		policyLogger = nil
	} else if fields != nil {
		policyLogger = policyLogger.WithFields(fields)
	} else {
		policyLogger = policyLogger.WithFields(logrus.Fields{
			logfields.LogSubsys:              subsystem,
			logfields.EndpointID:             e.ID,
			logfields.ContainerID:            e.getShortContainerID(),
			logfields.DatapathPolicyRevision: e.policyRevision,
			logfields.DesiredPolicyRevision:  e.nextPolicyRevision,
			logfields.IPv4:                   e.IPv4.String(),
			logfields.IPv6:                   e.IPv6.String(),
			logfields.K8sPodName:             e.getK8sNamespaceAndPodName(),
		})

		if e.SecurityIdentity != nil {
			policyLogger = policyLogger.WithField(logfields.Identity, e.SecurityIdentity.ID.StringID())
		}
	}
	atomic.StorePointer(&e.policyLogger, unsafe.Pointer(policyLogger))
}
<commit_msg>pkg/endpoint: fix data race in endpoint logger<commit_after>// SPDX-License-Identifier: Apache-2.0
// Copyright Authors of Cilium

package endpoint

import (
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"sync/atomic"
	"unsafe"

	"github.com/sirupsen/logrus"
	"gopkg.in/natefinch/lumberjack.v2"

	"github.com/cilium/cilium/pkg/logging"
	"github.com/cilium/cilium/pkg/logging/logfields"
	"github.com/cilium/cilium/pkg/option"
)

var (
	log       = logging.DefaultLogger.WithField(logfields.LogSubsys, subsystem)
	policyLog = logrus.New()

	policyLogOnce sync.Once
)

const (
	subsystem = "endpoint"

	fieldRegenLevel = "regeneration-level"
)

// getLogger returns a logrus object with EndpointID, containerID and the Endpoint
// revision fields.
func (e *Endpoint) getLogger() *logrus.Entry {
	v := atomic.LoadPointer(&e.logger)
	return (*logrus.Entry)(v)
}

// getPolicyLogger returns a logger to be used for policy update debugging, or nil,
// if not configured.
func (e *Endpoint) getPolicyLogger() *logrus.Entry {
	v := atomic.LoadPointer(&e.policyLogger)
	return (*logrus.Entry)(v)
}

// PolicyDebug logs the 'msg' with 'fields' if policy debug logging is enabled.
func (e *Endpoint) PolicyDebug(fields logrus.Fields, msg string) {
	if dbgLog := e.getPolicyLogger(); dbgLog != nil {
		dbgLog.WithFields(fields).Debug(msg)
	}
}

// Logger returns a logrus object with EndpointID, containerID and the Endpoint
// revision fields. The caller must specify their subsystem.
func (e *Endpoint) Logger(subsystem string) *logrus.Entry {
	if e == nil {
		return log.WithField(logfields.LogSubsys, subsystem)
	}

	return e.getLogger().WithField(logfields.LogSubsys, subsystem)
}

// UpdateLogger creates a logger instance specific to this endpoint. It will
// create a custom Debug logger for this endpoint when the option on it is set.
// If fields is not nil only the those specific fields will be updated in the
// endpoint's logger, otherwise a full update of those fields is executed.
//
// Note: You must hold Endpoint.mutex.Lock() to synchronize logger pointer
// updates if the endpoint is already exposed. Callers that create new
// endopoints do not need locks to call this.
func (e *Endpoint) UpdateLogger(fields map[string]interface{}) {
	e.updatePolicyLogger(fields)
	v := atomic.LoadPointer(&e.logger)
	epLogger := (*logrus.Entry)(v)
	if fields != nil && epLogger != nil {
		newLogger := epLogger.WithFields(fields)
		atomic.StorePointer(&e.logger, unsafe.Pointer(newLogger))
		return
	}

	// We need to update if
	// - e.logger is nil (this happens on the first ever call to UpdateLogger via
	//   Logger above). This clause has to come first to guard the others.
	// - If any of EndpointID, containerID or policyRevision are different on the
	//   endpoint from the logger.
	// - The debug option on the endpoint is true, and the logger is not debug,
	//   or vice versa.
	shouldUpdate := epLogger == nil || (e.Options != nil &&
		e.Options.IsEnabled(option.Debug) != (epLogger.Level == logrus.DebugLevel))

	// do nothing if we do not need an update
	if !shouldUpdate {
		return
	}

	// default to a new default logger
	baseLogger := logging.InitializeDefaultLogger()

	// If this endpoint is set to debug ensure it will print debug by giving it
	// an independent logger
	if e.Options != nil && e.Options.IsEnabled(option.Debug) {
		baseLogger.SetLevel(logrus.DebugLevel)
	} else {
		// Debug mode takes priority; if not in debug, check what log level user
		// has set and set the endpoint's log to log at that level.
		baseLogger.SetLevel(logging.DefaultLogger.Level)
	}

	// When adding new fields, make sure they are abstracted by a setter
	// and update the logger when the value is set.
	l := baseLogger.WithFields(logrus.Fields{
		logfields.LogSubsys:              subsystem,
		logfields.EndpointID:             e.ID,
		logfields.ContainerID:            e.getShortContainerID(),
		logfields.DatapathPolicyRevision: e.policyRevision,
		logfields.DesiredPolicyRevision:  e.nextPolicyRevision,
		logfields.IPv4:                   e.IPv4.String(),
		logfields.IPv6:                   e.IPv6.String(),
		logfields.K8sPodName:             e.getK8sNamespaceAndPodName(),
	})

	if e.SecurityIdentity != nil {
		l = l.WithField(logfields.Identity, e.SecurityIdentity.ID.StringID())
	}

	atomic.StorePointer(&e.logger, unsafe.Pointer(l))
}

// Only to be called from UpdateLogger() above
func (e *Endpoint) updatePolicyLogger(fields map[string]interface{}) {
	pv := atomic.LoadPointer(&e.policyLogger)
	policyLogger := (*logrus.Entry)(pv)
	// e.Options check needed for unit testing.
	if policyLogger == nil && e.Options != nil && e.Options.IsEnabled(option.DebugPolicy) {
		policyLogOnce.Do(func() {
			maxSize := 10 // 10 MB
			if ms := os.Getenv("CILIUM_DBG_POLICY_LOG_MAX_SIZE"); ms != "" {
				if ms, err := strconv.Atoi(ms); err == nil {
					maxSize = ms
				}
			}
			maxBackups := 3
			if mb := os.Getenv("CILIUM_DBG_POLICY_LOG_MAX_BACKUPS"); mb != "" {
				if mb, err := strconv.Atoi(mb); err == nil {
					maxBackups = mb
				}
			}
			lumberjackLogger := &lumberjack.Logger{
				Filename:   filepath.Join(option.Config.StateDir, "endpoint-policy.log"),
				MaxSize:    maxSize,
				MaxBackups: maxBackups,
				MaxAge:     28, // days
				LocalTime:  true,
				Compress:   true,
			}
			policyLog.SetOutput(lumberjackLogger)
			policyLog.SetLevel(logrus.DebugLevel)
		})
		policyLogger = logrus.NewEntry(policyLog)
	}
	if policyLogger == nil || e.Options == nil {
		return
	}

	if !e.Options.IsEnabled(option.DebugPolicy) {
		policyLogger = nil
	} else if fields != nil {
		policyLogger = policyLogger.WithFields(fields)
	} else {
		policyLogger = policyLogger.WithFields(logrus.Fields{
			logfields.LogSubsys:              subsystem,
			logfields.EndpointID:             e.ID,
			logfields.ContainerID:            e.getShortContainerID(),
			logfields.DatapathPolicyRevision: e.policyRevision,
			logfields.DesiredPolicyRevision:  e.nextPolicyRevision,
			logfields.IPv4:                   e.IPv4.String(),
			logfields.IPv6:                   e.IPv6.String(),
			logfields.K8sPodName:             e.getK8sNamespaceAndPodName(),
		})

		if e.SecurityIdentity != nil {
			policyLogger = policyLogger.WithField(logfields.Identity, e.SecurityIdentity.ID.StringID())
		}
	}
	atomic.StorePointer(&e.policyLogger, unsafe.Pointer(policyLogger))
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/go-ini/ini"
	_ "github.com/go-sql-driver/mysql"
)

// HTTPResponse Structure used to define response object of every route request
type HTTPResponse struct {
	Status  bool   `json:"status"`
	Content string `json:"content"`
}

// ResponseType Constant
const ResponseType = "application/json"

// ContentType Constant
const ContentType = "Content-Type"

var db *sql.DB
var lag int

func main() {

	var portstring string

	flag.StringVar(&portstring, "port", "3307", "Listening port")
	flag.Parse()

	cfg, err := ini.Load(os.Getenv("HOME") + "/.my.cnf")

	if err != nil {
		log.Panic(err)
	}

	dbUser := cfg.Section("client").Key("user").String()
	dbPass := cfg.Section("client").Key("password").String()
	dbHost := cfg.Section("client").Key("hostname").String()

	db, err = sql.Open("mysql", dbUser+":"+dbPass+"@"+dbHost+"/mysql")

	if err := db.Ping(); err != nil {
		log.Panic(err)
	}

	defer db.Close()

	router := http.NewServeMux()

	router.HandleFunc("/status/ro", RouteStatusReadOnly)
	router.HandleFunc("/status/rw", RouteStatusReadWritable)
	router.HandleFunc("/status/single", RouteStatusSingle)
	router.HandleFunc("/status/leader", RouteStatusLeader)
	router.HandleFunc("/status/follower", RouteStatusFollower)
	router.HandleFunc("/status/topology", RouteStatusTopology)

	router.HandleFunc("/role/master", RouteRoleMaster)
	router.HandleFunc("/role/replica", RouteRoleReplica)
	router.HandleFunc("/role/replica/", RouteRoleReplicaByLag)
	router.HandleFunc("/role/galera", RouteRoleGalera)

	router.HandleFunc("/read/galera/state", RouteReadGaleraState)
	router.HandleFunc("/read/replication/lag", RouteReadReplicationLag)
	router.HandleFunc("/read/replication/master", RouteReadReplicationMaster)
	router.HandleFunc("/read/replication/replicas_count", RouteReadReplicasCounter)

	log.Printf("Listening on port %s ...", portstring)

	err2 := http.ListenAndServe(":"+portstring, LogRequests(CheckURL(router)))
	log.Fatal(err2)
}

/*
 *	Middleware layers
 */

// LogRequests Middleware level to log API requests
func LogRequests(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()

		next.ServeHTTP(w, r)

		log.Printf(
			"[%s]\t%s\t%s",
			r.Method,
			r.URL.String(),
			time.Since(start),
		)
	})
}

// CheckURL Middleware level to validate requested URI
func CheckURL(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		path := r.URL.String()
		pathLength := len(path)
		matchPath := "/role/replica/"
		matchLength := len(matchPath)

		if strings.Contains(path, matchPath) && pathLength > matchLength {
			lag, _ = strconv.Atoi(strings.Trim(path, matchPath))
		} else if strings.Compare(path, strings.TrimRight(path, "/")) != 0 {
			return
		}

		w.Header().Set(ContentType, ResponseType)

		next.ServeHTTP(w, r)
	})
}

/*
 *	General functions
 */

// int2bool Convert integers to boolean
func int2bool(value int) bool {
	if value != 0 {
		return true
	}

	return false
}

// unknownColumns Used to get value from specific column of a range of unknown columns
func unknownColumns(rows *sql.Rows, column string) string {
	columns, _ := rows.Columns()
	count := len(columns)
	values := make([]interface{}, count)
	valuePtrs := make([]interface{}, count)

	for rows.Next() {
		for i := range columns {
			valuePtrs[i] = &values[i]
		}

		rows.Scan(valuePtrs...)

		for i, col := range columns {

			var value interface{}

			val := values[i]

			b, ok := val.([]byte)

			if ok {
				value = string(b)
			} else {
				value = val
			}

			sNum := value.(string)

			if col == column {
				return sNum
			}
		}
	}

	return ""
}

// routeResponse Used to build response to API requests
func routeResponse(w http.ResponseWriter, httpStatus bool, contents string) {
	res := new(HTTPResponse)

	if httpStatus {
		w.WriteHeader(200)
	} else {
		w.WriteHeader(403)
	}

	res.Status = httpStatus
	res.Content = contents
	response, _ := json.Marshal(res)
	fmt.Fprintf(w, "%s", response)
}

/*
 *	Database functions
 */

// readOnly Check if database is in readonly mode, or not
func readOnly() bool {
	var state string
	var key string

	err := db.QueryRow("show variables like 'read_only'").Scan(&key, &state)

	if state == "OFF" || err != nil {
		return false
	}

	return true
}

// replicaStatus Read database status if it is a replica
func replicaStatus(lagCount int) (bool, int) {
	if lagCount == 0 {
		lagCount = 1<<63 - 1
	}

	rows, err := db.Query("show slave status")

	if err != nil {
		return false, 0
	}

	defer rows.Close()

	secondsBehindMaster := unknownColumns(rows, "Seconds_Behind_Master")

	if secondsBehindMaster == "" {
		secondsBehindMaster = "0"
	}

	lag, _ = strconv.Atoi(secondsBehindMaster)

	if lag > 0 {
		if lagCount > lag {
			return true, lag
		}

		return false, lag
	}

	return false, 0
}

// isReplica Get database's master, in case it is a replica
func isReplica() (bool, string) {
	rows, err := db.Query("show slave status")

	if err != nil {
		return false, ""
	}

	defer rows.Close()

	masterHost := unknownColumns(rows, "Master_Host")

	if masterHost != "" {
		return true, masterHost
	}

	return false, ""
}

// servingBinlogs ...
func servingBinlogs() int {
	var count int

	err := db.QueryRow(
		"select count(*) as n " +
			"from information_schema.processlist " +
			"where command = 'Binlog Dump'").Scan(&count)

	if err != nil {
		return 0
	}

	return count
}

// galeraClusterState ...
func galeraClusterState() (bool, string) {
	var v string

	err := db.QueryRow(
		"select variable_value as v " +
			"from information_schema.global_status " +
			"where variable_name like 'wsrep_local_state' = 4").Scan(&v)

	if err == sql.ErrNoRows || err != nil {
		return false, ""
	}

	return true, v
}

/*
 * Status routes
 */

// RouteStatusReadOnly ...
func RouteStatusReadOnly(w http.ResponseWriter, r *http.Request) {
	log.Print("Checking database status: readOnly...")
	isReadonly := readOnly()

	routeResponse(w, isReadonly, "")
}

// RouteStatusReadWritable ...
func RouteStatusReadWritable(w http.ResponseWriter, r *http.Request) {
	log.Print("Checking database status: readable and writable...")
	isReadonly := readOnly()

	routeResponse(w, !isReadonly, "")
}

// RouteStatusSingle ...
func RouteStatusSingle(w http.ResponseWriter, r *http.Request) {
	log.Print("Checking database status: single...")
	isReadonly := readOnly()
	isReplica, _ := isReplica()
	isServeLogs := int2bool(servingBinlogs())

	routeResponse(w, !isReadonly && !isReplica && !isServeLogs, "")
}

// RouteStatusLeader ...
func RouteStatusLeader(w http.ResponseWriter, r *http.Request) {
	log.Print("Checking database status: leader...")
	isReplica, _ := isReplica()
	isServeLogs := int2bool(servingBinlogs())

	routeResponse(w, !isReplica && isServeLogs, "")
}

// RouteStatusFollower ...
func RouteStatusFollower(w http.ResponseWriter, r *http.Request) {
	log.Print("Checking database status: follower...")
	isReplica, _ := isReplica()

	routeResponse(w, isReplica, "")

}

// RouteStatusTopology ...
func RouteStatusTopology(w http.ResponseWriter, r *http.Request) {
	log.Print("Checking database status: topology...")
	isReplica, _ := isReplica()
	replicaStatus, _ := replicaStatus(0)
	isServeLogs := int2bool(servingBinlogs())

	routeResponse(w, (!replicaStatus && isServeLogs) || isReplica, "")
}

/*
 * Roles routes
 */

// RouteRoleMaster ...
func RouteRoleMaster(w http.ResponseWriter, r *http.Request) {
	log.Print("Checking database role: master...")
	isReadonly := readOnly()
	isReplica, _ := isReplica()

	routeResponse(w, !isReadonly && !isReplica, "")
}

// RouteRoleReplica ...
func RouteRoleReplica(w http.ResponseWriter, r *http.Request) {
	log.Print("Checking database role: replica...")
	isReadonly := readOnly()
	replicaStatus, _ := replicaStatus(0)

	routeResponse(w, isReadonly && replicaStatus, "")
}

// RouteRoleReplicaByLag ...
func RouteRoleReplicaByLag(w http.ResponseWriter, r *http.Request) {
	log.Print("Checking database role: replica by lag...")
	isReadonly := readOnly()
	replicaStatus, _ := replicaStatus(lag)

	routeResponse(w, isReadonly && replicaStatus, "")
}

// RouteRoleGalera ...
func RouteRoleGalera(w http.ResponseWriter, r *http.Request) {
	log.Print("Checking database role: galera...")
	galeraClusterState, _ := galeraClusterState()

	routeResponse(w, galeraClusterState, "")
}

/*
 * Read routes
 */

// RouteReadGaleraState ...
func RouteReadGaleraState(w http.ResponseWriter, r *http.Request) {
	log.Print("Reading database state: galera...")
	galeraClusterState, varValue := galeraClusterState()

	routeResponse(w, galeraClusterState, varValue)
}

// RouteReadReplicationLag ...
func RouteReadReplicationLag(w http.ResponseWriter, r *http.Request) {
	log.Print("Reading database replication: lag...")
	lagString := ""
	isReplica, _ := isReplica()
	_, lagValue := replicaStatus(0)

	if isReplica {
		lagString = strconv.Itoa(lagValue)
	}

	routeResponse(w, isReplica, lagString)
}

// RouteReadReplicationMaster ...
func RouteReadReplicationMaster(w http.ResponseWriter, r *http.Request) {
	log.Print("Reading database status: master...")
	lagString := ""
	isReplica, _ := isReplica()
	_, lagValue := replicaStatus(0)

	if isReplica {
		lagString = strconv.Itoa(lagValue)
	}

	routeResponse(w, isReplica, lagString)
}

// RouteReadReplicasCounter ...
func RouteReadReplicasCounter(w http.ResponseWriter, r *http.Request) {
	log.Print("Reading counter of database replications...")
	lagString := "0"
	isServeLogs := servingBinlogs()

	if int2bool(isServeLogs) {
		lagString = strconv.Itoa(isServeLogs)
	}

	routeResponse(w, int2bool(isServeLogs), lagString)

}
<commit_msg>Changes on MySQL API<commit_after>package main

import (
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	"github.com/go-ini/ini"
	_ "github.com/go-sql-driver/mysql"
)

// HTTPResponse Structure used to define response object of every route request
type HTTPResponse struct {
	Status  bool   `json:"status"`
	Content string `json:"content"`
}

// ResponseType Constant
const ResponseType = "application/json"

// ContentType Constant
const ContentType = "Content-Type"

var db *sql.DB
var lag int

func main() {

	var portstring string

	flag.StringVar(&portstring, "port", "3307", "Listening port")
	flag.Parse()

	cfg, err := ini.Load(os.Getenv("HOME") + "/.my.cnf")

	if err != nil {
		log.Panic(err)
	}

	var dbHost string

	dbUser := cfg.Section("client").Key("user").String()
	dbPass := cfg.Section("client").Key("password").String()

	isSocket := cfg.Section("client").HasKey("socket")

	if isSocket {
		dbHost = "unix(" + cfg.Section("client").Key("socket").String() + ")"
	} else {
		dbHost = cfg.Section("client").Key("hostname").String()
	}

	db, err = sql.Open("mysql", dbUser+":"+dbPass+"@"+dbHost+"/mysql")

	if err := db.Ping(); err != nil {
		log.Panic(err)
	}

	defer db.Close()

	router := http.NewServeMux()

	router.HandleFunc("/status/ro", RouteStatusReadOnly)
	router.HandleFunc("/status/rw", RouteStatusReadWritable)
	router.HandleFunc("/status/single", RouteStatusSingle)
	router.HandleFunc("/status/leader", RouteStatusLeader)
	router.HandleFunc("/status/follower", RouteStatusFollower)
	router.HandleFunc("/status/topology", RouteStatusTopology)

	router.HandleFunc("/role/master", RouteRoleMaster)
	router.HandleFunc("/role/replica", RouteRoleReplica)
	router.HandleFunc("/role/replica/", RouteRoleReplicaByLag)
	router.HandleFunc("/role/galera", RouteRoleGalera)

	router.HandleFunc("/read/galera/state", RouteReadGaleraState)
	router.HandleFunc("/read/replication/lag", RouteReadReplicationLag)
	router.HandleFunc("/read/replication/master", RouteReadReplicationMaster)
	router.HandleFunc("/read/replication/replicas_count", RouteReadReplicasCounter)

	log.Printf("Listening on port %s ...", portstring)

	err2 := http.ListenAndServe(":"+portstring, LogRequests(CheckURL(router)))
	log.Fatal(err2)
}

/*
 *	Middleware layers
 */

// LogRequests Middleware level to log API requests
func LogRequests(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()

		next.ServeHTTP(w, r)

		log.Printf(
			"[%s]\t%s\t%s",
			r.Method,
			r.URL.String(),
			time.Since(start),
		)
	})
}

// CheckURL Middleware level to validate requested URI
func CheckURL(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		path := r.URL.String()
		pathLength := len(path)
		matchPath := "/role/replica/"
		matchLength := len(matchPath)

		if strings.Contains(path, matchPath) && pathLength > matchLength {
			lag, _ = strconv.Atoi(strings.Trim(path, matchPath))
		} else if strings.Compare(path, strings.TrimRight(path, "/")) != 0 {
			return
		}

		w.Header().Set(ContentType, ResponseType)

		next.ServeHTTP(w, r)
	})
}

/*
 *	General functions
 */

// int2bool Convert integers to boolean
func int2bool(value int) bool {
	if value != 0 {
		return true
	}

	return false
}

// unknownColumns Used to get value from specific column of a range of unknown columns
func unknownColumns(rows *sql.Rows, column string) string {
	columns, _ := rows.Columns()
	count := len(columns)
	values := make([]interface{}, count)
	valuePtrs := make([]interface{}, count)

	for rows.Next() {
		for i := range columns {
			valuePtrs[i] = &values[i]
		}

		rows.Scan(valuePtrs...)

		for i, col := range columns {

			var value interface{}

			val := values[i]

			b, ok := val.([]byte)

			if ok {
				value = string(b)
			} else {
				value = val
			}

			sNum := value.(string)

			if col == column {
				return sNum
			}
		}
	}

	return ""
}

// routeResponse Used to build response to API requests
func routeResponse(w http.ResponseWriter, httpStatus bool, contents string) {
	res := new(HTTPResponse)

	if httpStatus {
		w.WriteHeader(200)
	} else {
		w.WriteHeader(403)
	}

	res.Status = httpStatus
	res.Content = contents
	response, _ := json.Marshal(res)
	fmt.Fprintf(w, "%s", response)
}

/*
 *	Database functions
 */

// readOnly Check if database is in readonly mode, or not
func readOnly() bool {
	var state string
	var key string

	err := db.QueryRow("show variables like 'read_only'").Scan(&key, &state)

	if state == "OFF" || err != nil {
		return false
	}

	return true
}

// replicaStatus Read database status if it is a replica
func replicaStatus(lagCount int) (bool, int) {
	if lagCount == 0 {
		lagCount = 1<<63 - 1
	}

	rows, err := db.Query("show slave status")

	if err != nil {
		return false, 0
	}

	defer rows.Close()

	secondsBehindMaster := unknownColumns(rows, "Seconds_Behind_Master")

	if secondsBehindMaster == "" {
		secondsBehindMaster = "0"
	}

	lag, _ = strconv.Atoi(secondsBehindMaster)

	if lag > 0 {
		if lagCount > lag {
			return true, lag
		}

		return false, lag
	}

	return false, 0
}

// isReplica Get database's master, in case it is a replica
func isReplica() (bool, string) {
	rows, err := db.Query("show slave status")

	if err != nil {
		return false, ""
	}

	defer rows.Close()

	masterHost := unknownColumns(rows, "Master_Host")

	if masterHost != "" {
		return true, masterHost
	}

	return false, ""
}

// servingBinlogs ...
func servingBinlogs() int {
	var count int

	err := db.QueryRow(
		"select count(*) as n " +
			"from information_schema.processlist " +
			"where command = 'Binlog Dump'").Scan(&count)

	if err != nil {
		return 0
	}

	return count
}

// galeraClusterState ...
func galeraClusterState() (bool, string) {
	var v string

	err := db.QueryRow(
		"select variable_value as v " +
			"from information_schema.global_status " +
			"where variable_name like 'wsrep_local_state' = 4").Scan(&v)

	if err == sql.ErrNoRows || err != nil {
		return false, ""
	}

	return true, v
}

/*
 * Status routes
 */

// RouteStatusReadOnly ...
func RouteStatusReadOnly(w http.ResponseWriter, r *http.Request) {
	log.Print("Checking database status: readOnly...")
	isReadonly := readOnly()

	routeResponse(w, isReadonly, "")
}

// RouteStatusReadWritable ...
func RouteStatusReadWritable(w http.ResponseWriter, r *http.Request) {
	log.Print("Checking database status: readable and writable...")
	isReadonly := readOnly()

	routeResponse(w, !isReadonly, "")
}

// RouteStatusSingle ...
func RouteStatusSingle(w http.ResponseWriter, r *http.Request) {
	log.Print("Checking database status: single...")
	isReadonly := readOnly()
	isReplica, _ := isReplica()
	isServeLogs := int2bool(servingBinlogs())

	routeResponse(w, !isReadonly && !isReplica && !isServeLogs, "")
}

// RouteStatusLeader ...
func RouteStatusLeader(w http.ResponseWriter, r *http.Request) {
	log.Print("Checking database status: leader...")
	isReplica, _ := isReplica()
	isServeLogs := int2bool(servingBinlogs())

	routeResponse(w, !isReplica && isServeLogs, "")
}

// RouteStatusFollower ...
func RouteStatusFollower(w http.ResponseWriter, r *http.Request) {
	log.Print("Checking database status: follower...")
	isReplica, _ := isReplica()

	routeResponse(w, isReplica, "")

}

// RouteStatusTopology ...
func RouteStatusTopology(w http.ResponseWriter, r *http.Request) {
	log.Print("Checking database status: topology...")
	isReplica, _ := isReplica()
	replicaStatus, _ := replicaStatus(0)
	isServeLogs := int2bool(servingBinlogs())

	routeResponse(w, (!replicaStatus && isServeLogs) || isReplica, "")
}

/*
 * Roles routes
 */

// RouteRoleMaster ...
func RouteRoleMaster(w http.ResponseWriter, r *http.Request) {
	log.Print("Checking database role: master...")
	isReadonly := readOnly()
	isReplica, _ := isReplica()

	routeResponse(w, !isReadonly && !isReplica, "")
}

// RouteRoleReplica ...
func RouteRoleReplica(w http.ResponseWriter, r *http.Request) {
	log.Print("Checking database role: replica...")
	isReadonly := readOnly()
	replicaStatus, _ := replicaStatus(0)

	routeResponse(w, isReadonly && replicaStatus, "")
}

// RouteRoleReplicaByLag ...
func RouteRoleReplicaByLag(w http.ResponseWriter, r *http.Request) {
	log.Print("Checking database role: replica by lag...")
	isReadonly := readOnly()
	replicaStatus, _ := replicaStatus(lag)

	routeResponse(w, isReadonly && replicaStatus, "")
}

// RouteRoleGalera ...
func RouteRoleGalera(w http.ResponseWriter, r *http.Request) {
	log.Print("Checking database role: galera...")
	galeraClusterState, _ := galeraClusterState()

	routeResponse(w, galeraClusterState, "")
}

/*
 * Read routes
 */

// RouteReadGaleraState ...
func RouteReadGaleraState(w http.ResponseWriter, r *http.Request) {
	log.Print("Reading database state: galera...")
	galeraClusterState, varValue := galeraClusterState()

	routeResponse(w, galeraClusterState, varValue)
}

// RouteReadReplicationLag ...
func RouteReadReplicationLag(w http.ResponseWriter, r *http.Request) {
	log.Print("Reading database replication: lag...")
	lagString := ""
	isReplica, _ := isReplica()
	_, lagValue := replicaStatus(0)

	if isReplica {
		lagString = strconv.Itoa(lagValue)
	}

	routeResponse(w, isReplica, lagString)
}

// RouteReadReplicationMaster ...
func RouteReadReplicationMaster(w http.ResponseWriter, r *http.Request) {
	log.Print("Reading database status: master...")
	lagString := ""
	isReplica, _ := isReplica()
	_, lagValue := replicaStatus(0)

	if isReplica {
		lagString = strconv.Itoa(lagValue)
	}

	routeResponse(w, isReplica, lagString)
}

// RouteReadReplicasCounter ...
func RouteReadReplicasCounter(w http.ResponseWriter, r *http.Request) {
	log.Print("Reading counter of database replications...")
	lagString := "0"
	isServeLogs := servingBinlogs()

	if int2bool(isServeLogs) {
		lagString = strconv.Itoa(isServeLogs)
	}

	routeResponse(w, int2bool(isServeLogs), lagString)

}
<|endoftext|>


<commit_before>package main

import (
    "github.com/dshearer/jobber"
    "net"
    "net/rpc"
    "fmt"
    "os"
    "os/user"
    "syscall"
    "flag"
)

const (
    ListCmdStr   = "list"
    LogCmdStr    = "log"
    ReloadCmdStr = "reload"
    StopCmdStr   = "stop"
    TestCmdStr   = "test"
)

func usage() {
    fmt.Printf("Usage: %v [flags] (%v|%v|%v|%v|%v)\nFlags:\n", os.Args[0], ListCmdStr, LogCmdStr, ReloadCmdStr, StopCmdStr, TestCmdStr)
    flag.PrintDefaults()
}

func subcmdUsage(subcmd string, flagSet *flag.FlagSet) func() {
    return func() {
        fmt.Printf("\nUsage: %v %v [flags]\nFlags:\n", os.Args[0], subcmd)
        flagSet.PrintDefaults()
    }
}

func failIfNotRoot(user *user.User) {
    if user.Uid != "0" {
        fmt.Fprintf(os.Stderr, "You must be root.\n")
        os.Exit(1)
    }
}

func main() {
    flag.Usage = usage

    var helpFlag_p = flag.Bool("h", false, "help")
    flag.Parse()

    if *helpFlag_p {
        usage()
        os.Exit(0)
    } else {
        if len(flag.Args()) == 0 {
            fmt.Fprintf(os.Stderr, "Please specify a command.\n\n")
            flag.Usage()
            os.Exit(1)
        }

        // make sure the daemon is running
        if _, err := os.Stat(jobber.DaemonSocketAddr); os.IsNotExist(err) {
            if flag.Arg(0) == StopCmdStr {
                os.Exit(0)
            } else {
                fmt.Fprintf(os.Stderr, "jobberd isn't running.\n")
                os.Exit(1)
            }
        }

        // connect to daemon
        addr, err := net.ResolveUnixAddr("unix", jobber.DaemonSocketAddr)
        if err != nil {
            fmt.Fprintf(os.Stderr, "Couldn't resolve Unix addr: %v\n", err)
            os.Exit(1)
        }
        conn, err := net.DialUnix("unix", nil, addr)
        if err != nil {
            fmt.Fprintf(os.Stderr, "Couldn't connect to daemon: %v\n", err)
            os.Exit(1)
        }
        defer conn.Close()
        rpcClient := rpc.NewClient(conn)
        if err != nil {
            fmt.Fprintf(os.Stderr, "Couldn't make RPC client: %v\n", err)
            os.Exit(1)
        }

        // drop privileges
        err = syscall.Setreuid(syscall.Getuid(), syscall.Getuid())
        if err != nil {
            fmt.Fprintf(os.Stderr, "Couldn't drop privileges: %v\n", err)
            os.Exit(1)
        }

        // get current username
        user, err := user.Current()
        if err != nil {
            fmt.Fprintf(os.Stderr, "Couldn't get current user: %v\n", err)
            os.Exit(1)
        }

        // do command
        switch flag.Arg(0) {
        case ListCmdStr:
            doListCmd(flag.Args()[1:], rpcClient, user)

        case LogCmdStr:
            doLogCmd(flag.Args()[1:], rpcClient, user)

        case ReloadCmdStr:
            doReloadCmd(flag.Args()[1:], rpcClient, user)

        case StopCmdStr:
            var result string
            arg := jobber.IpcArg{User: user.Username, ForAllUsers: false}
            err = rpcClient.Call("RealIpcServer.Stop", arg, &result)
            if err != nil {
                fmt.Fprintf(os.Stderr, "%v\n", err)
                os.Exit(1)
            }

        case TestCmdStr:
            doTestCmd(flag.Args()[1:], rpcClient, user)

        default:
            fmt.Fprintf(os.Stderr, "Invalid command: \"%v\".\n", flag.Arg(0))
            flag.Usage()
            os.Exit(1)
        }
    }
}

func doListCmd(args []string, rpcClient *rpc.Client, user *user.User) {
    // parse flags
    flagSet := flag.NewFlagSet("list", flag.ExitOnError)
    flagSet.Usage = subcmdUsage("list", flagSet)
    var help_p = flagSet.Bool("h", false, "help")
    var allUsers_p = flagSet.Bool("a", false, "all-users")
    flagSet.Parse(args)

    if *help_p {
        flagSet.Usage()
        os.Exit(0)
    } else {
        var result string
        arg := jobber.IpcArg{User: user.Username, ForAllUsers: *allUsers_p}
        err := rpcClient.Call("RealIpcServer.ListJobs", arg, &result)
        if err != nil {
            fmt.Fprintf(os.Stderr, "%v\n", err)
            os.Exit(1)
        }

        // print result
        fmt.Printf("%s\n", result)
    }
}

func doLogCmd(args []string, rpcClient *rpc.Client, user *user.User) {
    // parse flags
    flagSet := flag.NewFlagSet("log", flag.ExitOnError)
    flagSet.Usage = subcmdUsage("log", flagSet)
    var help_p = flagSet.Bool("h", false, "help")
    var allUsers_p = flagSet.Bool("a", false, "all-users")
    flagSet.Parse(args)

    if *help_p {
        flagSet.Usage()
        os.Exit(0)
    } else {
        var result string
        arg := jobber.IpcArg{User: user.Username, ForAllUsers: *allUsers_p}
        err := rpcClient.Call("RealIpcServer.ListHistory", arg, &result)
        if err != nil {
            fmt.Fprintf(os.Stderr, "%v\n", err)
            os.Exit(1)
        }

        // print result
        fmt.Printf("%s\n", result)
    }
}

func doReloadCmd(args []string, rpcClient *rpc.Client, user *user.User) {
    // parse flags
    flagSet := flag.NewFlagSet("reload", flag.ExitOnError)
    flagSet.Usage = subcmdUsage("reload", flagSet)
    var help_p = flagSet.Bool("h", false, "help")
    var allUsers_p = flagSet.Bool("a", false, "all-users")
    flagSet.Parse(args)

    if *help_p {
        flagSet.Usage()
        os.Exit(0)
    } else {
        var result string
        arg := jobber.IpcArg{User: user.Username, ForAllUsers: *allUsers_p}
        err := rpcClient.Call("RealIpcServer.Reload", arg, &result)
        if err != nil {
            fmt.Fprintf(os.Stderr, "%v\n", err)
            os.Exit(1)
        }
        fmt.Printf("%v\n", result)
    }
}

func doTestCmd(args []string, rpcClient *rpc.Client, user *user.User) {
    // parse flags
    flagSet := flag.NewFlagSet("test", flag.ExitOnError)
    flagSet.Usage = subcmdUsage("test", flagSet)
    var help_p *bool = flagSet.Bool("h", false, "help")
    var jobUser_p *string = flagSet.String("u", user.Username, "user")
    flagSet.Parse(args)

    if *help_p {
        flagSet.Usage()
        os.Exit(0)
    } else {
        // get job to test
        if len(flagSet.Args()) == 0 {
            fmt.Fprintf(os.Stderr, "You must specify a job to test.\n")
            os.Exit(1)
        }
        var job string = flagSet.Args()[0]

        // check "-u" opt
        if *jobUser_p == "" {
            fmt.Fprintf(os.Stderr, "Option requires an argument: \"-u\"\n")
            os.Exit(1)
        }

        var result string
        fmt.Printf("Running job \"%v\" for user \"%v\"...\n", job, *jobUser_p)
        arg := jobber.IpcArg{User: user.Username, Job: job, JobUser: *jobUser_p}
        err := rpcClient.Call("RealIpcServer.Test", arg, &result)
        if err != nil {
            fmt.Fprintf(os.Stderr, "%v\n", err)
            os.Exit(1)
        }
        fmt.Printf("%v\n", result)
    }
}
<commit_msg>Added '-v' option to show version.<commit_after>package main

import (
    "github.com/dshearer/jobber"
    "net"
    "net/rpc"
    "fmt"
    "os"
    "os/user"
    "syscall"
    "flag"
)

const (
    ListCmdStr   = "list"
    LogCmdStr    = "log"
    ReloadCmdStr = "reload"
    StopCmdStr   = "stop"
    TestCmdStr   = "test"
    Version      = "0.1.0"
)

func usage() {
    fmt.Printf("Usage: %v [flags] (%v|%v|%v|%v|%v)\nFlags:\n", os.Args[0], ListCmdStr, LogCmdStr, ReloadCmdStr, StopCmdStr, TestCmdStr)
    flag.PrintDefaults()
}

func version() {
    fmt.Printf("jobber %v\n", Version)
}

func subcmdUsage(subcmd string, flagSet *flag.FlagSet) func() {
    return func() {
        fmt.Printf("\nUsage: %v %v [flags]\nFlags:\n", os.Args[0], subcmd)
        flagSet.PrintDefaults()
    }
}

func failIfNotRoot(user *user.User) {
    if user.Uid != "0" {
        fmt.Fprintf(os.Stderr, "You must be root.\n")
        os.Exit(1)
    }
}

func main() {
    flag.Usage = usage

    var helpFlag_p = flag.Bool("h", false, "help")
    var versionFlag_p = flag.Bool("v", false, "version")
    flag.Parse()

    if *helpFlag_p {
        usage()
        os.Exit(0)
    } else if *versionFlag_p {
        version()
        os.Exit(0)
    } else {
        if len(flag.Args()) == 0 {
            fmt.Fprintf(os.Stderr, "Please specify a command.\n\n")
            flag.Usage()
            os.Exit(1)
        }

        // make sure the daemon is running
        if _, err := os.Stat(jobber.DaemonSocketAddr); os.IsNotExist(err) {
            if flag.Arg(0) == StopCmdStr {
                os.Exit(0)
            } else {
                fmt.Fprintf(os.Stderr, "jobberd isn't running.\n")
                os.Exit(1)
            }
        }

        // connect to daemon
        addr, err := net.ResolveUnixAddr("unix", jobber.DaemonSocketAddr)
        if err != nil {
            fmt.Fprintf(os.Stderr, "Couldn't resolve Unix addr: %v\n", err)
            os.Exit(1)
        }
        conn, err := net.DialUnix("unix", nil, addr)
        if err != nil {
            fmt.Fprintf(os.Stderr, "Couldn't connect to daemon: %v\n", err)
            os.Exit(1)
        }
        defer conn.Close()
        rpcClient := rpc.NewClient(conn)
        if err != nil {
            fmt.Fprintf(os.Stderr, "Couldn't make RPC client: %v\n", err)
            os.Exit(1)
        }

        // drop privileges
        err = syscall.Setreuid(syscall.Getuid(), syscall.Getuid())
        if err != nil {
            fmt.Fprintf(os.Stderr, "Couldn't drop privileges: %v\n", err)
            os.Exit(1)
        }

        // get current username
        user, err := user.Current()
        if err != nil {
            fmt.Fprintf(os.Stderr, "Couldn't get current user: %v\n", err)
            os.Exit(1)
        }

        // do command
        switch flag.Arg(0) {
        case ListCmdStr:
            doListCmd(flag.Args()[1:], rpcClient, user)

        case LogCmdStr:
            doLogCmd(flag.Args()[1:], rpcClient, user)

        case ReloadCmdStr:
            doReloadCmd(flag.Args()[1:], rpcClient, user)

        case StopCmdStr:
            var result string
            arg := jobber.IpcArg{User: user.Username, ForAllUsers: false}
            err = rpcClient.Call("RealIpcServer.Stop", arg, &result)
            if err != nil {
                fmt.Fprintf(os.Stderr, "%v\n", err)
                os.Exit(1)
            }

        case TestCmdStr:
            doTestCmd(flag.Args()[1:], rpcClient, user)

        default:
            fmt.Fprintf(os.Stderr, "Invalid command: \"%v\".\n", flag.Arg(0))
            flag.Usage()
            os.Exit(1)
        }
    }
}

func doListCmd(args []string, rpcClient *rpc.Client, user *user.User) {
    // parse flags
    flagSet := flag.NewFlagSet("list", flag.ExitOnError)
    flagSet.Usage = subcmdUsage("list", flagSet)
    var help_p = flagSet.Bool("h", false, "help")
    var allUsers_p = flagSet.Bool("a", false, "all-users")
    flagSet.Parse(args)

    if *help_p {
        flagSet.Usage()
        os.Exit(0)
    } else {
        var result string
        arg := jobber.IpcArg{User: user.Username, ForAllUsers: *allUsers_p}
        err := rpcClient.Call("RealIpcServer.ListJobs", arg, &result)
        if err != nil {
            fmt.Fprintf(os.Stderr, "%v\n", err)
            os.Exit(1)
        }

        // print result
        fmt.Printf("%s\n", result)
    }
}

func doLogCmd(args []string, rpcClient *rpc.Client, user *user.User) {
    // parse flags
    flagSet := flag.NewFlagSet("log", flag.ExitOnError)
    flagSet.Usage = subcmdUsage("log", flagSet)
    var help_p = flagSet.Bool("h", false, "help")
    var allUsers_p = flagSet.Bool("a", false, "all-users")
    flagSet.Parse(args)

    if *help_p {
        flagSet.Usage()
        os.Exit(0)
    } else {
        var result string
        arg := jobber.IpcArg{User: user.Username, ForAllUsers: *allUsers_p}
        err := rpcClient.Call("RealIpcServer.ListHistory", arg, &result)
        if err != nil {
            fmt.Fprintf(os.Stderr, "%v\n", err)
            os.Exit(1)
        }

        // print result
        fmt.Printf("%s\n", result)
    }
}

func doReloadCmd(args []string, rpcClient *rpc.Client, user *user.User) {
    // parse flags
    flagSet := flag.NewFlagSet("reload", flag.ExitOnError)
    flagSet.Usage = subcmdUsage("reload", flagSet)
    var help_p = flagSet.Bool("h", false, "help")
    var allUsers_p = flagSet.Bool("a", false, "all-users")
    flagSet.Parse(args)

    if *help_p {
        flagSet.Usage()
        os.Exit(0)
    } else {
        var result string
        arg := jobber.IpcArg{User: user.Username, ForAllUsers: *allUsers_p}
        err := rpcClient.Call("RealIpcServer.Reload", arg, &result)
        if err != nil {
            fmt.Fprintf(os.Stderr, "%v\n", err)
            os.Exit(1)
        }
        fmt.Printf("%v\n", result)
    }
}

func doTestCmd(args []string, rpcClient *rpc.Client, user *user.User) {
    // parse flags
    flagSet := flag.NewFlagSet("test", flag.ExitOnError)
    flagSet.Usage = subcmdUsage("test", flagSet)
    var help_p *bool = flagSet.Bool("h", false, "help")
    var jobUser_p *string = flagSet.String("u", user.Username, "user")
    flagSet.Parse(args)

    if *help_p {
        flagSet.Usage()
        os.Exit(0)
    } else {
        // get job to test
        if len(flagSet.Args()) == 0 {
            fmt.Fprintf(os.Stderr, "You must specify a job to test.\n")
            os.Exit(1)
        }
        var job string = flagSet.Args()[0]

        // check "-u" opt
        if *jobUser_p == "" {
            fmt.Fprintf(os.Stderr, "Option requires an argument: \"-u\"\n")
            os.Exit(1)
        }

        var result string
        fmt.Printf("Running job \"%v\" for user \"%v\"...\n", job, *jobUser_p)
        arg := jobber.IpcArg{User: user.Username, Job: job, JobUser: *jobUser_p}
        err := rpcClient.Call("RealIpcServer.Test", arg, &result)
        if err != nil {
            fmt.Fprintf(os.Stderr, "%v\n", err)
            os.Exit(1)
        }
        fmt.Printf("%v\n", result)
    }
}
<|endoftext|>


<commit_before>package escape // import "github.com/influxdata/influxdb/pkg/escape"

import "bytes"

func Bytes(in []byte) []byte {
	for b, esc := range Codes {
		in = bytes.Replace(in, []byte{b}, esc, -1)
	}
	return in
}

func Unescape(in []byte) []byte {
	i := 0
	inLen := len(in)
	var out []byte

	for {
		if i >= inLen {
			break
		}
		if in[i] == '\\' && i+1 < inLen {
			switch in[i+1] {
			case ',':
				out = append(out, ',')
				i += 2
				continue
			case '"':
				out = append(out, '"')
				i += 2
				continue
			case ' ':
				out = append(out, ' ')
				i += 2
				continue
			case '=':
				out = append(out, '=')
				i += 2
				continue
			}
		}
		out = append(out, in[i])
		i += 1
	}
	return out
}
<commit_msg>Avoid allocations when no escape chars present<commit_after>package escape // import "github.com/influxdata/influxdb/pkg/escape"

import "bytes"

func Bytes(in []byte) []byte {
	for b, esc := range Codes {
		in = bytes.Replace(in, []byte{b}, esc, -1)
	}
	return in
}

func Unescape(in []byte) []byte {
	if len(in) == 0 {
		return nil
	}

	var hasEscape bool
	for _, b := range in {
		if b == '\\' {
			hasEscape = true
			break
		}
	}

	if !hasEscape {
		return in
	}

	i := 0
	inLen := len(in)
	var out []byte

	for {
		if i >= inLen {
			break
		}
		if in[i] == '\\' && i+1 < inLen {
			switch in[i+1] {
			case ',':
				out = append(out, ',')
				i += 2
				continue
			case '"':
				out = append(out, '"')
				i += 2
				continue
			case ' ':
				out = append(out, ' ')
				i += 2
				continue
			case '=':
				out = append(out, '=')
				i += 2
				continue
			}
		}
		out = append(out, in[i])
		i += 1
	}
	return out
}
<|endoftext|>


<commit_before>// Copyright 2014 Wandoujia Inc. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package models

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"path"

	"github.com/ngaut/zkhelper"

	"github.com/juju/errors"
	"github.com/ngaut/go-zookeeper/zk"
	log "github.com/ngaut/logging"
)

const (
	PROXY_STATE_ONLINE       = "online"
	PROXY_STATE_OFFLINE      = "offline"
	PROXY_STATE_MARK_OFFLINE = "mark_offline"
)

type ProxyInfo struct {
	Id           string `json:"id"`
	Addr         string `json:"addr"`
	LastEvent    string `json:"last_event"`
	LastEventTs  int64  `json:"last_event_ts"`
	State        string `json:"state"`
	Description  string `json:"description"`
	DebugVarAddr string `json:"debug_var_addr"`
}

func (p ProxyInfo) Ops() (int64, error) {
	resp, err := http.Get("http://" + p.DebugVarAddr + "/debug/vars")
	if err != nil {
		return -1, errors.Trace(err)
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return -1, errors.Trace(err)
	}

	m := make(map[string]interface{})
	err = json.Unmarshal(body, &m)
	if err != nil {
		return -1, errors.Trace(err)
	}

	if v, ok := m["router"]; ok {
		if vv, ok := v.(map[string]interface{})["ops"]; ok {
			return int64(vv.(float64)), nil
		}
	}

	return 0, nil
}

func (p ProxyInfo) DebugVars() (map[string]interface{}, error) {
	resp, err := http.Get("http://" + p.DebugVarAddr + "/debug/vars")
	if err != nil {
		return nil, errors.Trace(err)
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, errors.Trace(err)
	}

	m := make(map[string]interface{})
	err = json.Unmarshal(body, &m)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return m, nil
}

func GetProxyPath(productName string) string {
	return fmt.Sprintf("/zk/codis/db_%s/proxy", productName)
}

func CreateProxyInfo(zkConn zkhelper.Conn, productName string, pi *ProxyInfo) (string, error) {
	data, err := json.Marshal(pi)
	if err != nil {
		return "", errors.Trace(err)
	}

	zkhelper.CreateRecursive(zkConn, GetProxyPath(productName), string(data),
		0, zkhelper.DefaultDirACLs())
	return zkConn.Create(path.Join(GetProxyPath(productName), pi.Id), data,
		zk.FlagEphemeral, zkhelper.DefaultDirACLs())
}

func ProxyList(zkConn zkhelper.Conn, productName string, filter func(*ProxyInfo) bool) ([]ProxyInfo, error) {
	var ret []ProxyInfo = make([]ProxyInfo, 0)
	root := GetProxyPath(productName)
	proxies, _, err := zkConn.Children(root)
	if err != nil && !zkhelper.ZkErrorEqual(err, zk.ErrNoNode) {
		return nil, errors.Trace(err)
	}

	for _, proxyName := range proxies {
		pi, err := GetProxyInfo(zkConn, productName, proxyName)
		if err != nil {
			return nil, errors.Trace(err)
		}
		if filter == nil || filter(pi) == true {
			ret = append(ret, *pi)
		}
	}

	return ret, nil
}

var ErrUnknownProxyStatus = errors.New("unknown status, should be (online offline)")

func SetProxyStatus(zkConn zkhelper.Conn, productName string, proxyName string, status string) error {
	p, err := GetProxyInfo(zkConn, productName, proxyName)
	if err != nil {
		return errors.Trace(err)
	}

	if status != PROXY_STATE_ONLINE && status != PROXY_STATE_MARK_OFFLINE && status != PROXY_STATE_OFFLINE {
		return errors.Errorf("%v, %s", ErrUnknownProxyStatus, status)
	}

	p.State = status
	b, _ := json.Marshal(p)

	_, err = zkConn.Set(path.Join(GetProxyPath(productName), proxyName), b, -1)
	if err != nil {
		return errors.Trace(err)
	}

	if status == PROXY_STATE_MARK_OFFLINE {
		// wait for the proxy down
		for {
			_, _, c, err := zkConn.GetW(path.Join(GetProxyPath(productName), proxyName))
			if zkhelper.ZkErrorEqual(err, zk.ErrNoNode) {
				return nil
			} else if err != nil {
				return errors.Trace(err)
			}
			<-c
			info, err := GetProxyInfo(zkConn, productName, proxyName)
			log.Info("mark_offline, check proxy status:", proxyName, info, err)
			if zkhelper.ZkErrorEqual(err, zk.ErrNoNode) {
				log.Info("shutdown proxy successful")
				return nil
			} else if err != nil {
				return errors.Trace(err)
			}
			if info.State == PROXY_STATE_OFFLINE {
				log.Info("proxy:", proxyName, "offline success!")
				return nil
			}
		}
	}

	return nil
}

func GetProxyInfo(zkConn zkhelper.Conn, productName string, proxyName string) (*ProxyInfo, error) {
	var pi ProxyInfo
	data, _, err := zkConn.Get(path.Join(GetProxyPath(productName), proxyName))
	if err != nil {
		return nil, errors.Trace(err)
	}

	if err := json.Unmarshal(data, &pi); err != nil {
		return nil, errors.Trace(err)
	}

	return &pi, nil
}
<commit_msg>clean up<commit_after>// Copyright 2014 Wandoujia Inc. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package models

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"path"

	"github.com/ngaut/zkhelper"

	"github.com/juju/errors"
	"github.com/ngaut/go-zookeeper/zk"
	log "github.com/ngaut/logging"
)

const (
	PROXY_STATE_ONLINE       = "online"
	PROXY_STATE_OFFLINE      = "offline"
	PROXY_STATE_MARK_OFFLINE = "mark_offline"
)

type ProxyInfo struct {
	Id           string `json:"id"`
	Addr         string `json:"addr"`
	LastEvent    string `json:"last_event"`
	LastEventTs  int64  `json:"last_event_ts"`
	State        string `json:"state"`
	Description  string `json:"description"`
	DebugVarAddr string `json:"debug_var_addr"`
}

func (p ProxyInfo) Ops() (int64, error) {
	resp, err := http.Get("http://" + p.DebugVarAddr + "/debug/vars")
	if err != nil {
		return -1, errors.Trace(err)
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return -1, errors.Trace(err)
	}

	m := make(map[string]interface{})
	err = json.Unmarshal(body, &m)
	if err != nil {
		return -1, errors.Trace(err)
	}

	if v, ok := m["router"]; ok {
		if vv, ok := v.(map[string]interface{})["ops"]; ok {
			return int64(vv.(float64)), nil
		}
	}

	return 0, nil
}

func (p ProxyInfo) DebugVars() (map[string]interface{}, error) {
	resp, err := http.Get("http://" + p.DebugVarAddr + "/debug/vars")
	if err != nil {
		return nil, errors.Trace(err)
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, errors.Trace(err)
	}

	m := make(map[string]interface{})
	err = json.Unmarshal(body, &m)
	if err != nil {
		return nil, errors.Trace(err)
	}
	return m, nil
}

func GetProxyPath(productName string) string {
	return fmt.Sprintf("/zk/codis/db_%s/proxy", productName)
}

func CreateProxyInfo(zkConn zkhelper.Conn, productName string, pi *ProxyInfo) (string, error) {
	data, err := json.Marshal(pi)
	if err != nil {
		return "", errors.Trace(err)
	}

	zkhelper.CreateRecursive(zkConn, GetProxyPath(productName), string(data),
		0, zkhelper.DefaultDirACLs())
	return zkConn.Create(path.Join(GetProxyPath(productName), pi.Id), data,
		zk.FlagEphemeral, zkhelper.DefaultDirACLs())
}

func ProxyList(zkConn zkhelper.Conn, productName string, filter func(*ProxyInfo) bool) ([]ProxyInfo, error) {
	ret := make([]ProxyInfo, 0)
	root := GetProxyPath(productName)
	proxies, _, err := zkConn.Children(root)
	if err != nil && !zkhelper.ZkErrorEqual(err, zk.ErrNoNode) {
		return nil, errors.Trace(err)
	}

	for _, proxyName := range proxies {
		pi, err := GetProxyInfo(zkConn, productName, proxyName)
		if err != nil {
			return nil, errors.Trace(err)
		}
		if filter == nil || filter(pi) == true {
			ret = append(ret, *pi)
		}
	}

	return ret, nil
}

var ErrUnknownProxyStatus = errors.New("unknown status, should be (online offline)")

func SetProxyStatus(zkConn zkhelper.Conn, productName string, proxyName string, status string) error {
	p, err := GetProxyInfo(zkConn, productName, proxyName)
	if err != nil {
		return errors.Trace(err)
	}

	if status != PROXY_STATE_ONLINE && status != PROXY_STATE_MARK_OFFLINE && status != PROXY_STATE_OFFLINE {
		return errors.Errorf("%v, %s", ErrUnknownProxyStatus, status)
	}

	p.State = status
	b, _ := json.Marshal(p)

	_, err = zkConn.Set(path.Join(GetProxyPath(productName), proxyName), b, -1)
	if err != nil {
		return errors.Trace(err)
	}

	if status == PROXY_STATE_MARK_OFFLINE {
		// wait for the proxy down
		for {
			_, _, c, err := zkConn.GetW(path.Join(GetProxyPath(productName), proxyName))
			if zkhelper.ZkErrorEqual(err, zk.ErrNoNode) {
				return nil
			} else if err != nil {
				return errors.Trace(err)
			}
			<-c
			info, err := GetProxyInfo(zkConn, productName, proxyName)
			log.Info("mark_offline, check proxy status:", proxyName, info, err)
			if zkhelper.ZkErrorEqual(err, zk.ErrNoNode) {
				log.Info("shutdown proxy successful")
				return nil
			} else if err != nil {
				return errors.Trace(err)
			}
			if info.State == PROXY_STATE_OFFLINE {
				log.Info("proxy:", proxyName, "offline success!")
				return nil
			}
		}
	}

	return nil
}

func GetProxyInfo(zkConn zkhelper.Conn, productName string, proxyName string) (*ProxyInfo, error) {
	var pi ProxyInfo
	data, _, err := zkConn.Get(path.Join(GetProxyPath(productName), proxyName))
	if err != nil {
		return nil, errors.Trace(err)
	}

	if err := json.Unmarshal(data, &pi); err != nil {
		return nil, errors.Trace(err)
	}

	return &pi, nil
}
<|endoftext|>


<commit_before>// Copyright 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.

// Package fs contains an HTTP file system that works with zip contents.
package fs

import (
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"encoding/hex"
	"encoding/pem"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/cozy/cozy-stack/pkg/config"
	"github.com/cozy/cozy-stack/pkg/filetype"
	"github.com/cozy/cozy-stack/pkg/logger"
	multierror "github.com/hashicorp/go-multierror"
)

var assetsClient = &http.Client{
	Timeout: 30 * time.Second,
}

var globalAssets sync.Map // {context:path -> *Asset}

const sumLen = 10

type AssetOption struct {
	Name     string `json:"name"`
	Context  string `json:"context"`
	URL      string `json:"url"`
	Shasum   string `json:"shasum"`
	IsCustom bool   `json:"is_custom,omitempty"`
}

// Asset holds unzipped read-only file contents and file metadata.
type Asset struct {
	AssetOption
	Etag        string `json:"etag"`
	NameWithSum string `json:"name_with_sum"`
	Mime        string `json:"mime"`

	zippedData   []byte
	zippedSize   string
	unzippedData []byte
	unzippedSize string
}

func (f *Asset) Size() string {
	return f.unzippedSize
}
func (f *Asset) Reader() *bytes.Reader {
	return bytes.NewReader(f.unzippedData)
}

func (f *Asset) GzipSize() string {
	return f.zippedSize
}
func (f *Asset) GzipReader() *bytes.Reader {
	return bytes.NewReader(f.zippedData)
}

// cacheTTL is time to live for the assets cache.
var cacheTTL = 30 * 24 * time.Hour

// Register registers zip contents data, later used to
// initialize the statik file system.
func Register(zipData string) {
	if zipData == "" {
		panic("statik/fs: no zip data registered")
	}
	if err := unzip([]byte(zipData)); err != nil {
		panic(fmt.Errorf("statik/fs: error unzipping data: %s", err))
	}
}

type Cache interface {
	Get(key string) (io.Reader, bool)
	Set(key string, data []byte, expiration time.Duration)
	RefreshTTL(key string, expiration time.Duration)
}

func RegisterCustomExternals(cache Cache, opts []AssetOption, maxTryCount int) error {
	if len(opts) == 0 {
		return nil
	}

	assetsCh := make(chan AssetOption)
	doneCh := make(chan error)

	for i := 0; i < len(opts); i++ {
		go func() {
			var err error
			sleepDuration := 500 * time.Millisecond
			opt := <-assetsCh

			for tryCount := 0; tryCount < maxTryCount+1; tryCount++ {
				err = registerCustomExternal(cache, opt)
				if err == nil {
					break
				}
				logger.WithNamespace("statik").
					Errorf("Could not load asset from %q, retrying in %s", opt.URL, sleepDuration)
				time.Sleep(sleepDuration)
				sleepDuration *= 4
			}

			doneCh <- err
		}()
	}

	for _, opt := range opts {
		assetsCh <- opt
	}
	close(assetsCh)

	var errm error
	for i := 0; i < len(opts); i++ {
		if err := <-doneCh; err != nil {
			errm = multierror.Append(errm, err)
		}
	}
	return errm
}

func registerCustomExternal(cache Cache, opt AssetOption) error {
	if opt.Context == "" {
		logger.WithNamespace("custom assets").
			Warningf("Could not load asset %s with empty context", opt.URL)
		return nil
	}

	name := normalizeAssetName(opt.Name)
	if currentAsset, ok := Get(name, opt.Context); ok {
		if currentAsset.Shasum == opt.Shasum {
			cache.RefreshTTL(name, cacheTTL)
			return nil
		}
	}

	opt.IsCustom = true

	assetURL := opt.URL

	var key string
	var body io.Reader
	var ok, storeInCache bool
	if opt.Shasum != "" {
		key = fmt.Sprintf("assets:%s:%s:%s", opt.Context, name, opt.Shasum)
		body, ok = cache.Get(key)
	}
	if !ok {
		u, err := url.Parse(assetURL)
		if err != nil {
			return err
		}

		switch u.Scheme {
		case "http", "https":
			req, err := http.NewRequest(http.MethodGet, assetURL, nil)
			if err != nil {
				return err
			}
			res, err := assetsClient.Do(req)
			if err != nil {
				return err
			}
			defer res.Body.Close()
			if res.StatusCode != http.StatusOK {
				return fmt.Errorf("could not load external asset on %s: status code %d", assetURL, res.StatusCode)
			}
			body = res.Body
		case "file":
			f, err := os.Open(u.Path)
			if err != nil {
				return err
			}
			defer f.Close()
			body = f
		default:
			return fmt.Errorf("does not support externals assets with scheme %q", u.Scheme)
		}

		storeInCache = true
	}

	h := sha256.New()
	zippedDataBuf := new(bytes.Buffer)
	gw := gzip.NewWriter(zippedDataBuf)

	teeReader := io.TeeReader(body, io.MultiWriter(h, gw))
	unzippedData, err := ioutil.ReadAll(teeReader)
	if err != nil {
		return err
	}
	if errc := gw.Close(); errc != nil {
		return errc
	}

	sum := h.Sum(nil)

	if opt.Shasum == "" {
		opt.Shasum = hex.EncodeToString(sum)
		key = fmt.Sprintf("assets:%s:%s:%s", opt.Context, name, opt.Shasum)
		log := logger.WithNamespace("custom_external")
		log.Warnf("shasum was not provided for file %s, inserting unsafe content %s: %s",
			opt.Name, opt.URL, opt.Shasum)
	}

	if hex.EncodeToString(sum) != opt.Shasum {
		return fmt.Errorf("external content checksum do not match: expected %s got %x on url %s",
			opt.Shasum, sum, assetURL)
	}

	if storeInCache {
		cache.Set(key, unzippedData, cacheTTL)
	}

	asset := newAsset(opt, zippedDataBuf.Bytes(), unzippedData)
	storeAsset(asset)
	return nil
}

func unzip(data []byte) (err error) {
	for {
		block, rest := pem.Decode(data)
		if block == nil {
			break
		}
		var zippedData, unzippedData []byte
		zippedData = block.Bytes
		var gr *gzip.Reader
		gr, err = gzip.NewReader(bytes.NewReader(block.Bytes))
		if err != nil {
			return
		}
		h := sha256.New()
		r := io.TeeReader(gr, h)
		unzippedData, err = ioutil.ReadAll(r)
		if err != nil {
			return
		}
		if err = gr.Close(); err != nil {
			return
		}

		name := block.Headers["Name"]
		opt := AssetOption{
			Name:    name,
			Context: config.DefaultInstanceContext,
			Shasum:  hex.EncodeToString(h.Sum(nil)),
		}
		asset := newAsset(opt, zippedData, unzippedData)
		storeAsset(asset)
		data = rest
	}
	return
}

func normalizeAssetName(name string) string {
	return path.Join("/", name)
}

func newAsset(opt AssetOption, zippedData, unzippedData []byte) *Asset {
	mime := filetype.ByExtension(path.Ext(opt.Name))
	if mime == "" {
		mime = filetype.Match(unzippedData)
	}

	sumx := opt.Shasum
	etag := fmt.Sprintf(`"%s"`, sumx[:sumLen])

	opt.Name = normalizeAssetName(opt.Name)

	nameWithSum := opt.Name
	nameBase := path.Base(opt.Name)
	if off := strings.IndexByte(nameBase, '.'); off >= 0 {
		nameDir := path.Dir(opt.Name)
		nameWithSum = path.Join("/", nameDir, nameBase[:off]+"."+sumx[:sumLen]+nameBase[off:])
	}

	return &Asset{
		AssetOption: opt,
		Etag:        etag,
		NameWithSum: nameWithSum,
		Mime:        mime,
		zippedData:  zippedData,
		zippedSize:  strconv.Itoa(len(zippedData)),

		unzippedData: unzippedData,
		unzippedSize: strconv.Itoa(len(unzippedData)),
	}
}

// threadsafe
func storeAsset(asset *Asset) {
	context := asset.Context
	if context == "" {
		context = config.DefaultInstanceContext
	}
	contextKey := marshalContextKey(context, asset.Name)
	globalAssets.Store(contextKey, asset)
}

func DeleteAsset(asset *Asset) {
	context := asset.Context
	if context == "" {
		context = config.DefaultInstanceContext
	}
	contextKey := marshalContextKey(context, asset.Name)
	globalAssets.Delete(contextKey)
}

func Get(name string, context ...string) (*Asset, bool) {
	var ctx string
	if len(context) > 0 && context[0] != "" {
		ctx = context[0]
	} else {
		ctx = config.DefaultInstanceContext
	}
	asset, ok := globalAssets.Load(marshalContextKey(ctx, name))
	if !ok {
		return nil, false
	}
	return asset.(*Asset), true
}

func Open(name string, context ...string) (*bytes.Reader, error) {
	f, ok := Get(name, context...)
	if ok {
		return f.Reader(), nil
	}
	return nil, os.ErrNotExist
}

func Foreach(predicate func(name, context string, f *Asset)) {
	globalAssets.Range(func(contextKey interface{}, v interface{}) bool {
		context, name, _ := unMarshalContextKey(contextKey.(string))
		predicate(name, context, v.(*Asset))
		return true
	})
}

func marshalContextKey(context, name string) (marshaledKey string) {
	return context + ":" + name
}

func unMarshalContextKey(contextKey string) (context string, name string, err error) {
	unmarshaled := strings.SplitN(contextKey, ":", 2)
	if len(unmarshaled) != 2 {
		panic("statik/fs: the contextKey is malformed")
	}
	return unmarshaled[0], unmarshaled[1], nil
}
<commit_msg>Fix refresh the TTL for the assets cache on read (#1813)<commit_after>// Copyright 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.

// Package fs contains an HTTP file system that works with zip contents.
package fs

import (
	"bytes"
	"compress/gzip"
	"crypto/sha256"
	"encoding/hex"
	"encoding/pem"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/cozy/cozy-stack/pkg/config"
	"github.com/cozy/cozy-stack/pkg/filetype"
	"github.com/cozy/cozy-stack/pkg/logger"
	multierror "github.com/hashicorp/go-multierror"
)

var assetsClient = &http.Client{
	Timeout: 30 * time.Second,
}

var globalAssets sync.Map // {context:path -> *Asset}

const sumLen = 10

type AssetOption struct {
	Name     string `json:"name"`
	Context  string `json:"context"`
	URL      string `json:"url"`
	Shasum   string `json:"shasum"`
	IsCustom bool   `json:"is_custom,omitempty"`
}

// Asset holds unzipped read-only file contents and file metadata.
type Asset struct {
	AssetOption
	Etag        string `json:"etag"`
	NameWithSum string `json:"name_with_sum"`
	Mime        string `json:"mime"`

	zippedData   []byte
	zippedSize   string
	unzippedData []byte
	unzippedSize string
}

func (f *Asset) Size() string {
	return f.unzippedSize
}
func (f *Asset) Reader() *bytes.Reader {
	return bytes.NewReader(f.unzippedData)
}

func (f *Asset) GzipSize() string {
	return f.zippedSize
}
func (f *Asset) GzipReader() *bytes.Reader {
	return bytes.NewReader(f.zippedData)
}

// cacheTTL is time to live for the assets cache.
var cacheTTL = 30 * 24 * time.Hour

// Register registers zip contents data, later used to
// initialize the statik file system.
func Register(zipData string) {
	if zipData == "" {
		panic("statik/fs: no zip data registered")
	}
	if err := unzip([]byte(zipData)); err != nil {
		panic(fmt.Errorf("statik/fs: error unzipping data: %s", err))
	}
}

type Cache interface {
	Get(key string) (io.Reader, bool)
	Set(key string, data []byte, expiration time.Duration)
	RefreshTTL(key string, expiration time.Duration)
}

func RegisterCustomExternals(cache Cache, opts []AssetOption, maxTryCount int) error {
	if len(opts) == 0 {
		return nil
	}

	assetsCh := make(chan AssetOption)
	doneCh := make(chan error)

	for i := 0; i < len(opts); i++ {
		go func() {
			var err error
			sleepDuration := 500 * time.Millisecond
			opt := <-assetsCh

			for tryCount := 0; tryCount < maxTryCount+1; tryCount++ {
				err = registerCustomExternal(cache, opt)
				if err == nil {
					break
				}
				logger.WithNamespace("statik").
					Errorf("Could not load asset from %q, retrying in %s", opt.URL, sleepDuration)
				time.Sleep(sleepDuration)
				sleepDuration *= 4
			}

			doneCh <- err
		}()
	}

	for _, opt := range opts {
		assetsCh <- opt
	}
	close(assetsCh)

	var errm error
	for i := 0; i < len(opts); i++ {
		if err := <-doneCh; err != nil {
			errm = multierror.Append(errm, err)
		}
	}
	return errm
}

func registerCustomExternal(cache Cache, opt AssetOption) error {
	if opt.Context == "" {
		logger.WithNamespace("custom assets").
			Warningf("Could not load asset %s with empty context", opt.URL)
		return nil
	}

	name := normalizeAssetName(opt.Name)
	if currentAsset, ok := Get(name, opt.Context); ok {
		if currentAsset.Shasum == opt.Shasum {
			return nil
		}
	}

	opt.IsCustom = true

	assetURL := opt.URL

	var key string
	var body io.Reader
	var ok, storeInCache bool
	if opt.Shasum != "" {
		key = fmt.Sprintf("assets:%s:%s:%s", opt.Context, name, opt.Shasum)
		body, ok = cache.Get(key)
	}
	if !ok {
		u, err := url.Parse(assetURL)
		if err != nil {
			return err
		}

		switch u.Scheme {
		case "http", "https":
			req, err := http.NewRequest(http.MethodGet, assetURL, nil)
			if err != nil {
				return err
			}
			res, err := assetsClient.Do(req)
			if err != nil {
				return err
			}
			defer res.Body.Close()
			if res.StatusCode != http.StatusOK {
				return fmt.Errorf("could not load external asset on %s: status code %d", assetURL, res.StatusCode)
			}
			body = res.Body
		case "file":
			f, err := os.Open(u.Path)
			if err != nil {
				return err
			}
			defer f.Close()
			body = f
		default:
			return fmt.Errorf("does not support externals assets with scheme %q", u.Scheme)
		}

		storeInCache = true
	}

	h := sha256.New()
	zippedDataBuf := new(bytes.Buffer)
	gw := gzip.NewWriter(zippedDataBuf)

	teeReader := io.TeeReader(body, io.MultiWriter(h, gw))
	unzippedData, err := ioutil.ReadAll(teeReader)
	if err != nil {
		return err
	}
	if errc := gw.Close(); errc != nil {
		return errc
	}

	sum := h.Sum(nil)

	if opt.Shasum == "" {
		opt.Shasum = hex.EncodeToString(sum)
		key = fmt.Sprintf("assets:%s:%s:%s", opt.Context, name, opt.Shasum)
		log := logger.WithNamespace("custom_external")
		log.Warnf("shasum was not provided for file %s, inserting unsafe content %s: %s",
			opt.Name, opt.URL, opt.Shasum)
	}

	if hex.EncodeToString(sum) != opt.Shasum {
		return fmt.Errorf("external content checksum do not match: expected %s got %x on url %s",
			opt.Shasum, sum, assetURL)
	}

	if storeInCache {
		cache.Set(key, unzippedData, cacheTTL)
	} else {
		cache.RefreshTTL(key, cacheTTL)
	}

	asset := newAsset(opt, zippedDataBuf.Bytes(), unzippedData)
	storeAsset(asset)
	return nil
}

func unzip(data []byte) (err error) {
	for {
		block, rest := pem.Decode(data)
		if block == nil {
			break
		}
		var zippedData, unzippedData []byte
		zippedData = block.Bytes
		var gr *gzip.Reader
		gr, err = gzip.NewReader(bytes.NewReader(block.Bytes))
		if err != nil {
			return
		}
		h := sha256.New()
		r := io.TeeReader(gr, h)
		unzippedData, err = ioutil.ReadAll(r)
		if err != nil {
			return
		}
		if err = gr.Close(); err != nil {
			return
		}

		name := block.Headers["Name"]
		opt := AssetOption{
			Name:    name,
			Context: config.DefaultInstanceContext,
			Shasum:  hex.EncodeToString(h.Sum(nil)),
		}
		asset := newAsset(opt, zippedData, unzippedData)
		storeAsset(asset)
		data = rest
	}
	return
}

func normalizeAssetName(name string) string {
	return path.Join("/", name)
}

func newAsset(opt AssetOption, zippedData, unzippedData []byte) *Asset {
	mime := filetype.ByExtension(path.Ext(opt.Name))
	if mime == "" {
		mime = filetype.Match(unzippedData)
	}

	sumx := opt.Shasum
	etag := fmt.Sprintf(`"%s"`, sumx[:sumLen])

	opt.Name = normalizeAssetName(opt.Name)

	nameWithSum := opt.Name
	nameBase := path.Base(opt.Name)
	if off := strings.IndexByte(nameBase, '.'); off >= 0 {
		nameDir := path.Dir(opt.Name)
		nameWithSum = path.Join("/", nameDir, nameBase[:off]+"."+sumx[:sumLen]+nameBase[off:])
	}

	return &Asset{
		AssetOption: opt,
		Etag:        etag,
		NameWithSum: nameWithSum,
		Mime:        mime,
		zippedData:  zippedData,
		zippedSize:  strconv.Itoa(len(zippedData)),

		unzippedData: unzippedData,
		unzippedSize: strconv.Itoa(len(unzippedData)),
	}
}

// threadsafe
func storeAsset(asset *Asset) {
	context := asset.Context
	if context == "" {
		context = config.DefaultInstanceContext
	}
	contextKey := marshalContextKey(context, asset.Name)
	globalAssets.Store(contextKey, asset)
}

func DeleteAsset(asset *Asset) {
	context := asset.Context
	if context == "" {
		context = config.DefaultInstanceContext
	}
	contextKey := marshalContextKey(context, asset.Name)
	globalAssets.Delete(contextKey)
}

func Get(name string, context ...string) (*Asset, bool) {
	var ctx string
	if len(context) > 0 && context[0] != "" {
		ctx = context[0]
	} else {
		ctx = config.DefaultInstanceContext
	}
	asset, ok := globalAssets.Load(marshalContextKey(ctx, name))
	if !ok {
		return nil, false
	}
	return asset.(*Asset), true
}

func Open(name string, context ...string) (*bytes.Reader, error) {
	f, ok := Get(name, context...)
	if ok {
		return f.Reader(), nil
	}
	return nil, os.ErrNotExist
}

func Foreach(predicate func(name, context string, f *Asset)) {
	globalAssets.Range(func(contextKey interface{}, v interface{}) bool {
		context, name, _ := unMarshalContextKey(contextKey.(string))
		predicate(name, context, v.(*Asset))
		return true
	})
}

func marshalContextKey(context, name string) (marshaledKey string) {
	return context + ":" + name
}

func unMarshalContextKey(contextKey string) (context string, name string, err error) {
	unmarshaled := strings.SplitN(contextKey, ":", 2)
	if len(unmarshaled) != 2 {
		panic("statik/fs: the contextKey is malformed")
	}
	return unmarshaled[0], unmarshaled[1], nil
}
<|endoftext|>