<commit_before>// The MIT License (MIT)
//
// Copyright (c) 2014 wandoulabs
// Copyright (c) 2014 siddontang
//
// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// Copyright 2015 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package server

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"time"
	// For pprof
	_ "net/http/pprof"

	"github.com/blacktear23/go-proxyprotocol"
	"github.com/juju/errors"
	"github.com/pingcap/tidb/config"
	"github.com/pingcap/tidb/metrics"
	"github.com/pingcap/tidb/mysql"
	"github.com/pingcap/tidb/sessionctx/variable"
	"github.com/pingcap/tidb/terror"
	"github.com/pingcap/tidb/util"
	log "github.com/sirupsen/logrus"
)

var (
	baseConnID uint32
)

var (
	errUnknownFieldType  = terror.ClassServer.New(codeUnknownFieldType, "unknown field type")
	errInvalidPayloadLen = terror.ClassServer.New(codeInvalidPayloadLen, "invalid payload length")
	errInvalidSequence   = terror.ClassServer.New(codeInvalidSequence, "invalid sequence")
	errInvalidType       = terror.ClassServer.New(codeInvalidType, "invalid type")
	errNotAllowedCommand = terror.ClassServer.New(codeNotAllowedCommand, "the used command is not allowed with this TiDB version")
	errAccessDenied      = terror.ClassServer.New(codeAccessDenied, mysql.MySQLErrName[mysql.ErrAccessDenied])
)

// DefaultCapability is the capability of the server when it is created using the default configuration.
// When server is configured with SSL, the server will have extra capabilities compared to DefaultCapability.
const defaultCapability = mysql.ClientLongPassword | mysql.ClientLongFlag |
	mysql.ClientConnectWithDB | mysql.ClientProtocol41 |
	mysql.ClientTransactions | mysql.ClientSecureConnection | mysql.ClientFoundRows |
	mysql.ClientMultiStatements | mysql.ClientMultiResults | mysql.ClientLocalFiles |
	mysql.ClientConnectAtts | mysql.ClientPluginAuth

// Server is the MySQL protocol server
type Server struct {
	cfg               *config.Config
	tlsConfig         *tls.Config
	driver            IDriver
	listener          net.Listener
	rwlock            *sync.RWMutex
	concurrentLimiter *TokenLimiter
	clients           map[uint32]*clientConn
	capability        uint32

	// stopListenerCh is used when a critical error occurred, we don't want to exit the process, because there may be
	// a supervisor automatically restart it, then new client connection will be created, but we can't server it.
	// So we just stop the listener and store to force clients to chose other TiDB servers.
	stopListenerCh chan struct{}
	statusServer   *http.Server
}

// ConnectionCount gets current connection count.
func (s *Server) ConnectionCount() int {
	var cnt int
	s.rwlock.RLock()
	cnt = len(s.clients)
	s.rwlock.RUnlock()
	return cnt
}

func (s *Server) getToken() *Token {
	return s.concurrentLimiter.Get()
}

func (s *Server) releaseToken(token *Token) {
	s.concurrentLimiter.Put(token)
}

// newConn creates a new *clientConn from a net.Conn.
// It allocates a connection ID and random salt data for authentication.
func (s *Server) newConn(conn net.Conn) *clientConn {
	cc := newClientConn(s)
	if s.cfg.Performance.TCPKeepAlive {
		if tcpConn, ok := conn.(*net.TCPConn); ok {
			if err := tcpConn.SetKeepAlive(true); err != nil {
				log.Error("failed to set tcp keep alive option:", err)
			}
		}
	}
	cc.setConn(conn)
	cc.salt = util.RandomBuf(20)
	return cc
}

func (s *Server) skipAuth() bool {
	return s.cfg.Socket != ""
}

// NewServer creates a new Server.
func NewServer(cfg *config.Config, driver IDriver) (*Server, error) {
	s := &Server{
		cfg:               cfg,
		driver:            driver,
		concurrentLimiter: NewTokenLimiter(cfg.TokenLimit),
		rwlock:            &sync.RWMutex{},
		clients:           make(map[uint32]*clientConn),
		stopListenerCh:    make(chan struct{}, 1),
	}
	s.loadTLSCertificates()

	s.capability = defaultCapability
	if s.tlsConfig != nil {
		s.capability |= mysql.ClientSSL
	}

	var err error
	if cfg.Socket != "" {
		if s.listener, err = net.Listen("unix", cfg.Socket); err == nil {
			log.Infof("Server is running MySQL Protocol through Socket [%s]", cfg.Socket)
		}
	} else {
		addr := fmt.Sprintf("%s:%d", s.cfg.Host, s.cfg.Port)
		if s.listener, err = net.Listen("tcp", addr); err == nil {
			log.Infof("Server is running MySQL Protocol at [%s]", addr)
		}
	}

	if cfg.ProxyProtocol.Networks != "" {
		pplistener, errProxy := proxyprotocol.NewListener(s.listener, cfg.ProxyProtocol.Networks,
			int(cfg.ProxyProtocol.HeaderTimeout))
		if errProxy != nil {
			log.Error("ProxyProtocol Networks parameter invalid")
			return nil, errors.Trace(errProxy)
		}
		log.Infof("Server is running MySQL Protocol (through PROXY Protocol) at [%s]", s.cfg.Host)
		s.listener = pplistener
	}

	if err != nil {
		return nil, errors.Trace(err)
	}

	// Init rand seed for randomBuf()
	rand.Seed(time.Now().UTC().UnixNano())
	return s, nil
}

func (s *Server) loadTLSCertificates() {
	defer func() {
		if s.tlsConfig != nil {
			log.Infof("Secure connection is enabled (client verification enabled = %v)", len(variable.SysVars["ssl_ca"].Value) > 0)
			variable.SysVars["have_openssl"].Value = "YES"
			variable.SysVars["have_ssl"].Value = "YES"
			variable.SysVars["ssl_cert"].Value = s.cfg.Security.SSLCert
			variable.SysVars["ssl_key"].Value = s.cfg.Security.SSLKey
		} else {
			log.Warn("Secure connection is NOT ENABLED")
		}
	}()

	if len(s.cfg.Security.SSLCert) == 0 || len(s.cfg.Security.SSLKey) == 0 {
		s.tlsConfig = nil
		return
	}

	tlsCert, err := tls.LoadX509KeyPair(s.cfg.Security.SSLCert, s.cfg.Security.SSLKey)
	if err != nil {
		log.Warn(errors.ErrorStack(err))
		s.tlsConfig = nil
		return
	}

	// Try loading CA cert.
	clientAuthPolicy := tls.NoClientCert
	var certPool *x509.CertPool
	if len(s.cfg.Security.SSLCA) > 0 {
		caCert, err := ioutil.ReadFile(s.cfg.Security.SSLCA)
		if err != nil {
			log.Warn(errors.ErrorStack(err))
		} else {
			certPool = x509.NewCertPool()
			if certPool.AppendCertsFromPEM(caCert) {
				clientAuthPolicy = tls.VerifyClientCertIfGiven
			}
			variable.SysVars["ssl_ca"].Value = s.cfg.Security.SSLCA
		}
	}
	s.tlsConfig = &tls.Config{
		Certificates: []tls.Certificate{tlsCert},
		ClientCAs:    certPool,
		ClientAuth:   clientAuthPolicy,
		MinVersion:   0,
	}
}

// Run runs the server.
func (s *Server) Run() error {
	metrics.ServerEventCounter.WithLabelValues(metrics.EventStart).Inc()

	// Start HTTP API to report tidb info such as TPS.
	if s.cfg.Status.ReportStatus {
		s.startStatusHTTP()
	}
	for {
		conn, err := s.listener.Accept()
		if err != nil {
			if opErr, ok := err.(*net.OpError); ok {
				if opErr.Err.Error() == "use of closed network connection" {
					return nil
				}
			}

			// If we got PROXY protocol error, we should continue accept.
			if proxyprotocol.IsProxyProtocolError(err) {
				log.Errorf("PROXY protocol error: %s", err.Error())
				continue
			}

			log.Errorf("accept error %s", err.Error())
			return errors.Trace(err)
		}
		if s.shouldStopListener() {
			err = conn.Close()
			terror.Log(errors.Trace(err))
			break
		}
		go s.onConn(conn)
	}
	err := s.listener.Close()
	terror.Log(errors.Trace(err))
	s.listener = nil
	for {
		metrics.ServerEventCounter.WithLabelValues(metrics.EventHang).Inc()
		log.Errorf("listener stopped, waiting for manual kill.")
		time.Sleep(time.Minute)
	}
}

func (s *Server) shouldStopListener() bool {
	select {
	case <-s.stopListenerCh:
		return true
	default:
		return false
	}
}

// Close closes the server.
func (s *Server) Close() {
	s.rwlock.Lock()
	defer s.rwlock.Unlock()

	if s.listener != nil {
		err := s.listener.Close()
		terror.Log(errors.Trace(err))
		s.listener = nil
	}
	if s.statusServer != nil {
		err := s.statusServer.Shutdown(nil)
		terror.Log(errors.Trace(err))
		s.statusServer = nil
	}
	metrics.ServerEventCounter.WithLabelValues(metrics.EventClose).Inc()
}

// onConn runs in its own goroutine, handles queries from this connection.
func (s *Server) onConn(c net.Conn) {
	conn := s.newConn(c)
	if err := conn.handshake(); err != nil {
		// Some keep alive services will send request to TiDB and disconnect immediately.
		// So we only record metrics.
		metrics.HandShakeErrorCounter.Inc()
		err = c.Close()
		terror.Log(errors.Trace(err))
		return
	}
	log.Infof("[con:%d] new connection %s", conn.connectionID, c.RemoteAddr().String())
	defer func() {
		log.Infof("[con:%d] close connection", conn.connectionID)
	}()
	s.rwlock.Lock()
	s.clients[conn.connectionID] = conn
	connections := len(s.clients)
	s.rwlock.Unlock()
	metrics.ConnGauge.Set(float64(connections))

	conn.Run()
}

// ShowProcessList implements the SessionManager interface.
func (s *Server) ShowProcessList() []util.ProcessInfo {
	var rs []util.ProcessInfo
	s.rwlock.RLock()
	for _, client := range s.clients {
		if atomic.LoadInt32(&client.status) == connStatusWaitShutdown {
			continue
		}
		rs = append(rs, client.ctx.ShowProcess())
	}
	s.rwlock.RUnlock()
	return rs
}

// Kill implements the SessionManager interface.
func (s *Server) Kill(connectionID uint64, query bool) {
	s.rwlock.Lock()
	defer s.rwlock.Unlock()
	log.Infof("[server] Kill connectionID %d, query %t]", connectionID, query)
	metrics.ServerEventCounter.WithLabelValues(metrics.EventKill).Inc()

	conn, ok := s.clients[uint32(connectionID)]
	if !ok {
		return
	}

	conn.mu.RLock()
	cancelFunc := conn.mu.cancelFunc
	conn.mu.RUnlock()
	if cancelFunc != nil {
		cancelFunc()
	}

	if !query {
		// Mark the client connection status as WaitShutdown, when the goroutine detect
		// this, it will end the dispatch loop and exit.
		atomic.StoreInt32(&conn.status, connStatusWaitShutdown)
	}
}

// GracefulDown waits all clients to close.
func (s *Server) GracefulDown() {
	log.Info("[server] graceful shutdown.")
	metrics.ServerEventCounter.WithLabelValues(metrics.EventGracefulDown).Inc()

	count := s.ConnectionCount()
	for i := 0; count > 0; i++ {
		time.Sleep(time.Second)
		s.kickIdleConnection()

		count = s.ConnectionCount()
		// Print information for every 30s.
		if i%30 == 0 {
			log.Infof("graceful shutdown...connection count %d\n", count)
		}
	}
}

func (s *Server) kickIdleConnection() {
	var conns []*clientConn
	s.rwlock.RLock()
	for _, cc := range s.clients {
		if cc.ShutdownOrNotify() {
			// Shutdowned conn will be closed by us, and notified conn will exist themselves.
			conns = append(conns, cc)
		}
	}
	s.rwlock.RUnlock()

	for _, cc := range conns {
		err := cc.Close()
		if err != nil {
			log.Error("close connection error:", err)
		}
	}
}

// Server error codes.
const (
	codeUnknownFieldType  = 1
	codeInvalidPayloadLen = 2
	codeInvalidSequence   = 3
	codeInvalidType       = 4

	codeNotAllowedCommand = 1148
	codeAccessDenied      = mysql.ErrAccessDenied
)

func init() {
	serverMySQLErrCodes := map[terror.ErrCode]uint16{
		codeNotAllowedCommand: mysql.ErrNotAllowedCommand,
		codeAccessDenied:      mysql.ErrAccessDenied,
	}
	terror.ErrClassToMySQLCodes[terror.ClassServer] = serverMySQLErrCodes
}
<commit_msg>server: fix CI nil pointer panic (#6733)<commit_after>// The MIT License (MIT)
//
// Copyright (c) 2014 wandoulabs
// Copyright (c) 2014 siddontang
//
// Permission is hereby granted, free of charge, to any person obtaining a copy of
// this software and associated documentation files (the "Software"), to deal in
// the Software without restriction, including without limitation the rights to
// use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
// the Software, and to permit persons to whom the Software is furnished to do so,
// subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.

// Copyright 2015 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.

package server

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net"
	"net/http"
	"sync"
	"sync/atomic"
	"time"
	// For pprof
	_ "net/http/pprof"

	"github.com/blacktear23/go-proxyprotocol"
	"github.com/juju/errors"
	"github.com/pingcap/tidb/config"
	"github.com/pingcap/tidb/metrics"
	"github.com/pingcap/tidb/mysql"
	"github.com/pingcap/tidb/sessionctx/variable"
	"github.com/pingcap/tidb/terror"
	"github.com/pingcap/tidb/util"
	log "github.com/sirupsen/logrus"
)

var (
	baseConnID uint32
)

var (
	errUnknownFieldType  = terror.ClassServer.New(codeUnknownFieldType, "unknown field type")
	errInvalidPayloadLen = terror.ClassServer.New(codeInvalidPayloadLen, "invalid payload length")
	errInvalidSequence   = terror.ClassServer.New(codeInvalidSequence, "invalid sequence")
	errInvalidType       = terror.ClassServer.New(codeInvalidType, "invalid type")
	errNotAllowedCommand = terror.ClassServer.New(codeNotAllowedCommand, "the used command is not allowed with this TiDB version")
	errAccessDenied      = terror.ClassServer.New(codeAccessDenied, mysql.MySQLErrName[mysql.ErrAccessDenied])
)

// DefaultCapability is the capability of the server when it is created using the default configuration.
// When server is configured with SSL, the server will have extra capabilities compared to DefaultCapability.
const defaultCapability = mysql.ClientLongPassword | mysql.ClientLongFlag |
	mysql.ClientConnectWithDB | mysql.ClientProtocol41 |
	mysql.ClientTransactions | mysql.ClientSecureConnection | mysql.ClientFoundRows |
	mysql.ClientMultiStatements | mysql.ClientMultiResults | mysql.ClientLocalFiles |
	mysql.ClientConnectAtts | mysql.ClientPluginAuth

// Server is the MySQL protocol server
type Server struct {
	cfg               *config.Config
	tlsConfig         *tls.Config
	driver            IDriver
	listener          net.Listener
	rwlock            *sync.RWMutex
	concurrentLimiter *TokenLimiter
	clients           map[uint32]*clientConn
	capability        uint32

	// stopListenerCh is used when a critical error occurred, we don't want to exit the process, because there may be
	// a supervisor automatically restart it, then new client connection will be created, but we can't server it.
	// So we just stop the listener and store to force clients to chose other TiDB servers.
	stopListenerCh chan struct{}
	statusServer   *http.Server
}

// ConnectionCount gets current connection count.
func (s *Server) ConnectionCount() int {
	var cnt int
	s.rwlock.RLock()
	cnt = len(s.clients)
	s.rwlock.RUnlock()
	return cnt
}

func (s *Server) getToken() *Token {
	return s.concurrentLimiter.Get()
}

func (s *Server) releaseToken(token *Token) {
	s.concurrentLimiter.Put(token)
}

// newConn creates a new *clientConn from a net.Conn.
// It allocates a connection ID and random salt data for authentication.
func (s *Server) newConn(conn net.Conn) *clientConn {
	cc := newClientConn(s)
	if s.cfg.Performance.TCPKeepAlive {
		if tcpConn, ok := conn.(*net.TCPConn); ok {
			if err := tcpConn.SetKeepAlive(true); err != nil {
				log.Error("failed to set tcp keep alive option:", err)
			}
		}
	}
	cc.setConn(conn)
	cc.salt = util.RandomBuf(20)
	return cc
}

func (s *Server) skipAuth() bool {
	return s.cfg.Socket != ""
}

// NewServer creates a new Server.
func NewServer(cfg *config.Config, driver IDriver) (*Server, error) {
	s := &Server{
		cfg:               cfg,
		driver:            driver,
		concurrentLimiter: NewTokenLimiter(cfg.TokenLimit),
		rwlock:            &sync.RWMutex{},
		clients:           make(map[uint32]*clientConn),
		stopListenerCh:    make(chan struct{}, 1),
	}
	s.loadTLSCertificates()

	s.capability = defaultCapability
	if s.tlsConfig != nil {
		s.capability |= mysql.ClientSSL
	}

	var err error
	if cfg.Socket != "" {
		if s.listener, err = net.Listen("unix", cfg.Socket); err == nil {
			log.Infof("Server is running MySQL Protocol through Socket [%s]", cfg.Socket)
		}
	} else {
		addr := fmt.Sprintf("%s:%d", s.cfg.Host, s.cfg.Port)
		if s.listener, err = net.Listen("tcp", addr); err == nil {
			log.Infof("Server is running MySQL Protocol at [%s]", addr)
		}
	}

	if cfg.ProxyProtocol.Networks != "" {
		pplistener, errProxy := proxyprotocol.NewListener(s.listener, cfg.ProxyProtocol.Networks,
			int(cfg.ProxyProtocol.HeaderTimeout))
		if errProxy != nil {
			log.Error("ProxyProtocol Networks parameter invalid")
			return nil, errors.Trace(errProxy)
		}
		log.Infof("Server is running MySQL Protocol (through PROXY Protocol) at [%s]", s.cfg.Host)
		s.listener = pplistener
	}

	if err != nil {
		return nil, errors.Trace(err)
	}

	// Init rand seed for randomBuf()
	rand.Seed(time.Now().UTC().UnixNano())
	return s, nil
}

func (s *Server) loadTLSCertificates() {
	defer func() {
		if s.tlsConfig != nil {
			log.Infof("Secure connection is enabled (client verification enabled = %v)", len(variable.SysVars["ssl_ca"].Value) > 0)
			variable.SysVars["have_openssl"].Value = "YES"
			variable.SysVars["have_ssl"].Value = "YES"
			variable.SysVars["ssl_cert"].Value = s.cfg.Security.SSLCert
			variable.SysVars["ssl_key"].Value = s.cfg.Security.SSLKey
		} else {
			log.Warn("Secure connection is NOT ENABLED")
		}
	}()

	if len(s.cfg.Security.SSLCert) == 0 || len(s.cfg.Security.SSLKey) == 0 {
		s.tlsConfig = nil
		return
	}

	tlsCert, err := tls.LoadX509KeyPair(s.cfg.Security.SSLCert, s.cfg.Security.SSLKey)
	if err != nil {
		log.Warn(errors.ErrorStack(err))
		s.tlsConfig = nil
		return
	}

	// Try loading CA cert.
	clientAuthPolicy := tls.NoClientCert
	var certPool *x509.CertPool
	if len(s.cfg.Security.SSLCA) > 0 {
		caCert, err := ioutil.ReadFile(s.cfg.Security.SSLCA)
		if err != nil {
			log.Warn(errors.ErrorStack(err))
		} else {
			certPool = x509.NewCertPool()
			if certPool.AppendCertsFromPEM(caCert) {
				clientAuthPolicy = tls.VerifyClientCertIfGiven
			}
			variable.SysVars["ssl_ca"].Value = s.cfg.Security.SSLCA
		}
	}
	s.tlsConfig = &tls.Config{
		Certificates: []tls.Certificate{tlsCert},
		ClientCAs:    certPool,
		ClientAuth:   clientAuthPolicy,
		MinVersion:   0,
	}
}

// Run runs the server.
func (s *Server) Run() error {
	metrics.ServerEventCounter.WithLabelValues(metrics.EventStart).Inc()

	// Start HTTP API to report tidb info such as TPS.
	if s.cfg.Status.ReportStatus {
		s.startStatusHTTP()
	}
	for {
		conn, err := s.listener.Accept()
		if err != nil {
			if opErr, ok := err.(*net.OpError); ok {
				if opErr.Err.Error() == "use of closed network connection" {
					return nil
				}
			}

			// If we got PROXY protocol error, we should continue accept.
			if proxyprotocol.IsProxyProtocolError(err) {
				log.Errorf("PROXY protocol error: %s", err.Error())
				continue
			}

			log.Errorf("accept error %s", err.Error())
			return errors.Trace(err)
		}
		if s.shouldStopListener() {
			err = conn.Close()
			terror.Log(errors.Trace(err))
			break
		}
		go s.onConn(conn)
	}
	err := s.listener.Close()
	terror.Log(errors.Trace(err))
	s.listener = nil
	for {
		metrics.ServerEventCounter.WithLabelValues(metrics.EventHang).Inc()
		log.Errorf("listener stopped, waiting for manual kill.")
		time.Sleep(time.Minute)
	}
}

func (s *Server) shouldStopListener() bool {
	select {
	case <-s.stopListenerCh:
		return true
	default:
		return false
	}
}

// Close closes the server.
func (s *Server) Close() {
	s.rwlock.Lock()
	defer s.rwlock.Unlock()

	if s.listener != nil {
		err := s.listener.Close()
		terror.Log(errors.Trace(err))
		s.listener = nil
	}
	if s.statusServer != nil {
		err := s.statusServer.Close()
		terror.Log(errors.Trace(err))
		s.statusServer = nil
	}
	metrics.ServerEventCounter.WithLabelValues(metrics.EventClose).Inc()
}

// onConn runs in its own goroutine, handles queries from this connection.
func (s *Server) onConn(c net.Conn) {
	conn := s.newConn(c)
	if err := conn.handshake(); err != nil {
		// Some keep alive services will send request to TiDB and disconnect immediately.
		// So we only record metrics.
		metrics.HandShakeErrorCounter.Inc()
		err = c.Close()
		terror.Log(errors.Trace(err))
		return
	}
	log.Infof("[con:%d] new connection %s", conn.connectionID, c.RemoteAddr().String())
	defer func() {
		log.Infof("[con:%d] close connection", conn.connectionID)
	}()
	s.rwlock.Lock()
	s.clients[conn.connectionID] = conn
	connections := len(s.clients)
	s.rwlock.Unlock()
	metrics.ConnGauge.Set(float64(connections))

	conn.Run()
}

// ShowProcessList implements the SessionManager interface.
func (s *Server) ShowProcessList() []util.ProcessInfo {
	var rs []util.ProcessInfo
	s.rwlock.RLock()
	for _, client := range s.clients {
		if atomic.LoadInt32(&client.status) == connStatusWaitShutdown {
			continue
		}
		rs = append(rs, client.ctx.ShowProcess())
	}
	s.rwlock.RUnlock()
	return rs
}

// Kill implements the SessionManager interface.
func (s *Server) Kill(connectionID uint64, query bool) {
	s.rwlock.Lock()
	defer s.rwlock.Unlock()
	log.Infof("[server] Kill connectionID %d, query %t]", connectionID, query)
	metrics.ServerEventCounter.WithLabelValues(metrics.EventKill).Inc()

	conn, ok := s.clients[uint32(connectionID)]
	if !ok {
		return
	}

	conn.mu.RLock()
	cancelFunc := conn.mu.cancelFunc
	conn.mu.RUnlock()
	if cancelFunc != nil {
		cancelFunc()
	}

	if !query {
		// Mark the client connection status as WaitShutdown, when the goroutine detect
		// this, it will end the dispatch loop and exit.
		atomic.StoreInt32(&conn.status, connStatusWaitShutdown)
	}
}

// GracefulDown waits all clients to close.
func (s *Server) GracefulDown() {
	log.Info("[server] graceful shutdown.")
	metrics.ServerEventCounter.WithLabelValues(metrics.EventGracefulDown).Inc()

	count := s.ConnectionCount()
	for i := 0; count > 0; i++ {
		time.Sleep(time.Second)
		s.kickIdleConnection()

		count = s.ConnectionCount()
		// Print information for every 30s.
		if i%30 == 0 {
			log.Infof("graceful shutdown...connection count %d\n", count)
		}
	}
}

func (s *Server) kickIdleConnection() {
	var conns []*clientConn
	s.rwlock.RLock()
	for _, cc := range s.clients {
		if cc.ShutdownOrNotify() {
			// Shutdowned conn will be closed by us, and notified conn will exist themselves.
			conns = append(conns, cc)
		}
	}
	s.rwlock.RUnlock()

	for _, cc := range conns {
		err := cc.Close()
		if err != nil {
			log.Error("close connection error:", err)
		}
	}
}

// Server error codes.
const (
	codeUnknownFieldType  = 1
	codeInvalidPayloadLen = 2
	codeInvalidSequence   = 3
	codeInvalidType       = 4

	codeNotAllowedCommand = 1148
	codeAccessDenied      = mysql.ErrAccessDenied
)

func init() {
	serverMySQLErrCodes := map[terror.ErrCode]uint16{
		codeNotAllowedCommand: mysql.ErrNotAllowedCommand,
		codeAccessDenied:      mysql.ErrAccessDenied,
	}
	terror.ErrClassToMySQLCodes[terror.ClassServer] = serverMySQLErrCodes
}
<|endoftext|>


<commit_before>// Package server provides I/O for Ios servers

package main

import (
	"bufio"
	"flag"
	"github.com/golang/glog"
	"github.com/heidi-ann/ios/cache"
	"github.com/heidi-ann/ios/config"
	"github.com/heidi-ann/ios/consensus"
	"github.com/heidi-ann/ios/msgs"
	"github.com/heidi-ann/ios/store"
	"io"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
)

var keyval *store.Store
var c *cache.Cache
var cons_io *msgs.Io

var notifyclient map[msgs.ClientRequest](chan msgs.ClientResponse)
var notifyclient_mutex sync.RWMutex

type Peer struct {
	id      int
	address string
	handled bool // TODO: replace with Mutex
}

var peers []Peer
var peers_mutex sync.RWMutex

var id = flag.Int("id", -1, "server ID")
var config_file = flag.String("config", "example.conf", "Server configuration file")
var disk_path = flag.String("disk", ".", "Path to directory to store persistent storage")

func openFile(filename string) (*bufio.Writer, *bufio.Reader, *os.File, bool) {
	// check if file exists already for logging
	var is_new bool
	if _, err := os.Stat(filename); os.IsNotExist(err) {
		glog.Info("Creating and opening file: ", filename)
		is_new = true
	} else {
		glog.Info("Opening file: ", filename)
		is_new = false
	}

	// open file
	file, err := os.OpenFile(filename, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0777)
	if err != nil {
		glog.Fatal(err)
	}

	// create writer and reader
	w := bufio.NewWriter(file)
	r := bufio.NewReader(file)
	return w, r, file, is_new
}

func stateMachine() {
	for {
		req := <-cons_io.OutgoingRequests
		glog.Info("Request has been safely replicated by consensus algorithm", req)

		// check if request already applied
		found, reply := c.Check(req)
		if found {
			glog.Info("Request found in cache and thus cannot be applied")
		} else {
			// apply request
			output := keyval.Process(req.Request)
			//keyval.Print()

			// write response to request cache
			reply = msgs.ClientResponse{
				req.ClientID, req.RequestID, output}
			c.Add(reply)
		}

		// if any handleRequests are waiting on this reply, then reply to them
		notifyclient_mutex.Lock()
		if notifyclient[req] != nil {
			notifyclient[req] <- reply
		}
		notifyclient_mutex.Unlock()
	}
}

func handleRequest(req msgs.ClientRequest) msgs.ClientResponse {
	glog.Info("Handling ", req.Request)

	// check if already applied
	found, res := c.Check(req)
	if found {
		glog.Info("Request found in cache")
		return res // FAST PASS
	}

	// CONSENESUS ALGORITHM HERE
	glog.Info("Passing request to consensus algorithm")
	cons_io.IncomingRequests <- req

	// wait for reply
	notifyclient_mutex.Lock()
	notifyclient[req] = make(chan msgs.ClientResponse)
	notifyclient_mutex.Unlock()
	reply := <-notifyclient[req]

	// check reply
	if reply.ClientID != req.ClientID {
		glog.Fatal("ClientID is different")
	}
	if reply.RequestID != req.RequestID {
		glog.Fatal("RequestID is different")
	}

	return reply
}

// iterative through peers and check there is a handler for each
// try to create one if not
func checkPeer() {
	for i := range peers {
		peers_mutex.RLock()
		failed := !peers[i].handled
		peers_mutex.RUnlock()
		if failed {
			glog.Info("Peer ", i, " is not currently connected")
			cn, err := net.Dial("tcp", peers[i].address)

			if err != nil {
				glog.Warning(err)
			} else {
				go handlePeer(cn, true)
			}
		} else {
			glog.Info("Peer ", i, " is currently connected")
		}
	}
}

func handlePeer(cn net.Conn, init bool) {
	addr := cn.RemoteAddr().String()
	if init {
		glog.Info("Outgoing peer connection to ", addr)
	} else {
		glog.Info("Incoming peer connection from ", addr)
	}

	defer glog.Warningf("Connection closed from %s ", addr)

	// handle requests
	reader := bufio.NewReader(cn)
	writer := bufio.NewWriter(cn)

	// exchange peer ID's
	_, _ = writer.WriteString(strconv.Itoa(*id) + "\n")
	_ = writer.Flush()
	text, _ := reader.ReadString('\n')
	glog.Info("Received ", text)
	peer_id, err := strconv.Atoi(strings.Trim(text, "\n"))
	if err != nil {
		glog.Warning(err)
		return
	}

	glog.Infof("Ready to handle traffic from peer %d at %s ", peer_id, addr)

	peers_mutex.Lock()
	peers[peer_id].handled = true
	peers_mutex.Unlock()

	close_err := make(chan error)
	go func() {
		for {
			// read request
			glog.Infof("Ready for next message from %d", peer_id)
			text, err := reader.ReadBytes(byte('\n'))
			if err != nil {
				glog.Warning(err)
				close_err <- err
				break
			}
			glog.Infof("Read from peer %d: ", peer_id, string(text))
			cons_io.Incoming.BytesToProtoMsg(text)

		}
	}()

	go func() {
		for {
			// send reply
			glog.Infof("Ready to send message to %d", peer_id)
			b, err := cons_io.OutgoingUnicast[peer_id].ProtoMsgToBytes()
			if err != nil {
				glog.Fatal("Could not marshal message")
			}
			glog.Infof("Sending to %d: %s", peer_id, string(b))
			_, err = writer.Write(b)
			_, err = writer.Write([]byte("\n"))
			if err != nil {
				glog.Warning(err)
				close_err <- err
				break
			}
			err = writer.Flush()
			if err != nil {
				glog.Warning(err)
				close_err <- err
				break
			}
			glog.Info("Sent")
		}
	}()

	// block until connection fails
	<-close_err

	// tidy up
	glog.Warningf("No longer able to handle traffic from peer %d at %s ", peer_id, addr)
	peers_mutex.Lock()
	peers[peer_id].handled = false
	peers_mutex.Unlock()
	cons_io.Failure <- peer_id
	cn.Close()
}

func handleConnection(cn net.Conn) {
	glog.Info("Incoming client connection from ",
		cn.RemoteAddr().String())

	reader := bufio.NewReader(cn)
	writer := bufio.NewWriter(cn)

	for {

		// read request
		glog.Info("Ready for Reading")
		text, err := reader.ReadBytes(byte('\n'))
		if err != nil {
			if err == io.EOF {
				break
			}
			glog.Warning(err)
			break
		}
		glog.Info("--------------------New request----------------------")
		glog.Info("Request: ", string(text))
		req := new(msgs.ClientRequest)
		err = msgs.Unmarshal(text, req)
		if err != nil {
			glog.Fatal(err)
		}

		// construct reply
		reply := handleRequest(*req)
		b, err := msgs.Marshal(reply)
		if err != nil {
			glog.Fatal("error:", err)
		}
		glog.Info(string(b))

		// send reply
		// TODO: FIX currently all server send back replies
		glog.Info("Sending ", string(b))
		n, err := writer.Write(b)
		if err != nil {
			glog.Fatal(err)
		}
		_, err = writer.Write([]byte("\n"))
		if err != nil {
			glog.Fatal(err)
		}

		// tidy up
		err = writer.Flush()
		glog.Info("Finished sending ", n, " bytes")

	}

	cn.Close()
}

func main() {
	// set up logging
	flag.Parse()
	defer glog.Flush()

	conf := config.ParseServerConfig(*config_file)
	if *id == -1 {
		glog.Fatal("ID is required")
	}

	glog.Info("Starting server ", *id)
	defer glog.Warning("Shutting down server ", *id)

	//set up state machine
	keyval = store.New()
	c = cache.Create()
	// setup IO
	cons_io = msgs.MakeIo(2000, len(conf.Peers.Address))

	notifyclient = make(map[msgs.ClientRequest](chan msgs.ClientResponse))
	notifyclient_mutex = sync.RWMutex{}
	go stateMachine()

	// setting up persistent log
	disk, disk_reader, disk_file, is_empty := openFile(*disk_path + "/persistent_log_" + strconv.Itoa(*id) + ".temp")
	defer disk.Flush()
	meta_disk, meta_disk_reader, meta_disk_file, is_new := openFile(*disk_path + "/persistent_data_" + strconv.Itoa(*id) + ".temp")
	defer meta_disk.Flush()

	// check persistent storage for commands
	found := false
	log := make([]msgs.Entry, 10000) // TODO: Remove hard coded limit
	log_length := 0

	if !is_empty {
		for {
			b, err := disk_reader.ReadBytes(byte('\n'))
			if err != nil {
				glog.Info("No more commands in persistent storage")
				break
			}
			found = true
			var update msgs.LogUpdate
			err = msgs.Unmarshal(b, &update)
			if err != nil {
				glog.Fatal("Cannot parse log update", err)
			}
			log[update.Index] = update.Entry
			if log_length < update.Index {
				log_length = update.Index
			}
			glog.Info("Adding from persistent storage :", update)
		}
	}
	log = log[:log_length]

	// check persistent storage for view
	view := 0
	if !is_new {
		for {
			b, err := meta_disk_reader.ReadBytes(byte('\n'))
			if err != nil {
				glog.Info("No more view updates in persistent storage")
				break
			}
			found = true
			view, _ = strconv.Atoi(string(b))
		}
	}

	// write updates to persistent storage
	go func() {
		for {
			view := <-cons_io.ViewPersist
			glog.Info("Updating view to ", view)
			_, err := meta_disk.Write([]byte(strconv.Itoa(view)))
			_, err = meta_disk.Write([]byte("\n"))
			if err != nil {
				glog.Fatal(err)
			}
			meta_disk_file.Sync()
			cons_io.ViewPersistFsync <- view
		}
	}()

	go func() {
		for {
			log := <-cons_io.LogPersist
			glog.Info("Updating log with ", log)
			b, err := msgs.Marshal(log)
			if err != nil {
				glog.Fatal(err)
			}
			// write to persistent storage
			n1, err := disk.Write(b)
			n2, err := disk.Write([]byte("\n"))
			if err != nil {
				glog.Fatal(err)
			}
			glog.Info(n1+n2, " bytes written to persistent log")
			disk_file.Sync()
			cons_io.LogPersistFsync <- log
		}
	}()

	// set up client server
	glog.Info("Starting up client server")
	client_port := strings.Split(conf.Clients.Address[*id],":")[1]
	listeningPort := ":" + client_port
	ln, err := net.Listen("tcp", listeningPort)
	if err != nil {
		glog.Fatal(err)
	}

	// handle for incoming clients
	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				glog.Fatal(err)
			}
			go handleConnection(conn)
		}
	}()

	//set up peer state
	peers = make([]Peer, len(conf.Peers.Address))
	for i := range conf.Peers.Address {
		peers[i] = Peer{
			i, conf.Peers.Address[i], false}
	}
	peers_mutex = sync.RWMutex{}

	//set up peer server
	glog.Info("Starting up peer server")
	peer_port := strings.Split(conf.Peers.Address[*id],":")[1]
	listeningPort = ":" + peer_port
	lnPeers, err := net.Listen("tcp", listeningPort)
	if err != nil {
		glog.Fatal(err)
	}

	// handle local peer (without sending network traffic)
	peers_mutex.Lock()
	peers[*id].handled = true
	peers_mutex.Unlock()
	from := &(cons_io.Incoming)
	go from.Forward(cons_io.OutgoingUnicast[*id])

	// handle for incoming peers
	go func() {
		for {
			conn, err := lnPeers.Accept()
			if err != nil {
				glog.Fatal(err)
			}
			go handlePeer(conn, false)
		}
	}()

	// regularly check if all peers are connected and retry if not
	go func() {
		for {
			checkPeer()
			time.Sleep(100 * time.Millisecond)
		}
	}()

	// setting up the consensus algorithm
	log_max_length := 1000
	if conf.Options.Length > 0 {
		log_max_length = conf.Options.Length
	}
	cons_config := consensus.Config{*id, len(conf.Peers.Address),
		log_max_length, conf.Options.BatchInterval, conf.Options.MaxBatch, conf.Options.DelegateReplication, conf.Options.WindowSize}
	if !found {
		glog.Info("Starting fresh consensus instance")
		go consensus.Init(cons_io, cons_config)
	} else {
		glog.Info("Restoring consensus instance")
		go consensus.Recover(cons_io, cons_config, view, log)
	}
	//go cons_io.DumpPersistentStorage()

	// tidy up
	glog.Info("Setup complete")

	// waiting for exit
	// always flush (whatever happens)
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
	sig := <-sigs
	disk.Flush()
	meta_disk.Flush()
	glog.Flush()
	glog.Warning("Shutting down due to ", sig)
}
<commit_msg>reducing logging<commit_after>// Package server provides I/O for Ios servers

package main

import (
	"bufio"
	"flag"
	"github.com/golang/glog"
	"github.com/heidi-ann/ios/cache"
	"github.com/heidi-ann/ios/config"
	"github.com/heidi-ann/ios/consensus"
	"github.com/heidi-ann/ios/msgs"
	"github.com/heidi-ann/ios/store"
	"io"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
)

var keyval *store.Store
var c *cache.Cache
var cons_io *msgs.Io

var notifyclient map[msgs.ClientRequest](chan msgs.ClientResponse)
var notifyclient_mutex sync.RWMutex

type Peer struct {
	id      int
	address string
	handled bool // TODO: replace with Mutex
}

var peers []Peer
var peers_mutex sync.RWMutex

var id = flag.Int("id", -1, "server ID")
var config_file = flag.String("config", "example.conf", "Server configuration file")
var disk_path = flag.String("disk", ".", "Path to directory to store persistent storage")

func openFile(filename string) (*bufio.Writer, *bufio.Reader, *os.File, bool) {
	// check if file exists already for logging
	var is_new bool
	if _, err := os.Stat(filename); os.IsNotExist(err) {
		glog.Info("Creating and opening file: ", filename)
		is_new = true
	} else {
		glog.Info("Opening file: ", filename)
		is_new = false
	}

	// open file
	file, err := os.OpenFile(filename, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0777)
	if err != nil {
		glog.Fatal(err)
	}

	// create writer and reader
	w := bufio.NewWriter(file)
	r := bufio.NewReader(file)
	return w, r, file, is_new
}

func stateMachine() {
	for {
		req := <-cons_io.OutgoingRequests
		glog.Info("Request has been safely replicated by consensus algorithm", req)

		// check if request already applied
		found, reply := c.Check(req)
		if found {
			glog.Info("Request found in cache and thus cannot be applied")
		} else {
			// apply request
			output := keyval.Process(req.Request)
			//keyval.Print()

			// write response to request cache
			reply = msgs.ClientResponse{
				req.ClientID, req.RequestID, output}
			c.Add(reply)
		}

		// if any handleRequests are waiting on this reply, then reply to them
		notifyclient_mutex.Lock()
		if notifyclient[req] != nil {
			notifyclient[req] <- reply
		}
		notifyclient_mutex.Unlock()
	}
}

func handleRequest(req msgs.ClientRequest) msgs.ClientResponse {
	glog.Info("Handling ", req.Request)

	// check if already applied
	found, res := c.Check(req)
	if found {
		glog.Info("Request found in cache")
		return res // FAST PASS
	}

	// CONSENESUS ALGORITHM HERE
	glog.Info("Passing request to consensus algorithm")
	cons_io.IncomingRequests <- req

	// wait for reply
	notifyclient_mutex.Lock()
	notifyclient[req] = make(chan msgs.ClientResponse)
	notifyclient_mutex.Unlock()
	reply := <-notifyclient[req]

	// check reply
	if reply.ClientID != req.ClientID {
		glog.Fatal("ClientID is different")
	}
	if reply.RequestID != req.RequestID {
		glog.Fatal("RequestID is different")
	}

	return reply
}

// iterative through peers and check there is a handler for each
// try to create one if not
func checkPeer() {
	for i := range peers {
		peers_mutex.RLock()
		failed := !peers[i].handled
		peers_mutex.RUnlock()
		if failed {
			glog.Info("Peer ", i, " is not currently connected")
			cn, err := net.Dial("tcp", peers[i].address)

			if err != nil {
				glog.Warning(err)
			} else {
				go handlePeer(cn, true)
			}
		} else {
			//glog.Info("Peer ", i, " is currently connected")
		}
	}
}

func handlePeer(cn net.Conn, init bool) {
	addr := cn.RemoteAddr().String()
	if init {
		glog.Info("Outgoing peer connection to ", addr)
	} else {
		glog.Info("Incoming peer connection from ", addr)
	}

	defer glog.Warningf("Connection closed from %s ", addr)

	// handle requests
	reader := bufio.NewReader(cn)
	writer := bufio.NewWriter(cn)

	// exchange peer ID's
	_, _ = writer.WriteString(strconv.Itoa(*id) + "\n")
	_ = writer.Flush()
	text, _ := reader.ReadString('\n')
	glog.Info("Received ", text)
	peer_id, err := strconv.Atoi(strings.Trim(text, "\n"))
	if err != nil {
		glog.Warning(err)
		return
	}

	glog.Infof("Ready to handle traffic from peer %d at %s ", peer_id, addr)

	peers_mutex.Lock()
	peers[peer_id].handled = true
	peers_mutex.Unlock()

	close_err := make(chan error)
	go func() {
		for {
			// read request
			glog.Infof("Ready for next message from %d", peer_id)
			text, err := reader.ReadBytes(byte('\n'))
			if err != nil {
				glog.Warning(err)
				close_err <- err
				break
			}
			glog.Infof("Read from peer %d: ", peer_id, string(text))
			cons_io.Incoming.BytesToProtoMsg(text)

		}
	}()

	go func() {
		for {
			// send reply
			glog.Infof("Ready to send message to %d", peer_id)
			b, err := cons_io.OutgoingUnicast[peer_id].ProtoMsgToBytes()
			if err != nil {
				glog.Fatal("Could not marshal message")
			}
			glog.Infof("Sending to %d: %s", peer_id, string(b))
			_, err = writer.Write(b)
			_, err = writer.Write([]byte("\n"))
			if err != nil {
				glog.Warning(err)
				close_err <- err
				break
			}
			err = writer.Flush()
			if err != nil {
				glog.Warning(err)
				close_err <- err
				break
			}
			glog.Info("Sent")
		}
	}()

	// block until connection fails
	<-close_err

	// tidy up
	glog.Warningf("No longer able to handle traffic from peer %d at %s ", peer_id, addr)
	peers_mutex.Lock()
	peers[peer_id].handled = false
	peers_mutex.Unlock()
	cons_io.Failure <- peer_id
	cn.Close()
}

func handleConnection(cn net.Conn) {
	glog.Info("Incoming client connection from ",
		cn.RemoteAddr().String())

	reader := bufio.NewReader(cn)
	writer := bufio.NewWriter(cn)

	for {

		// read request
		glog.Info("Ready for Reading")
		text, err := reader.ReadBytes(byte('\n'))
		if err != nil {
			if err == io.EOF {
				break
			}
			glog.Warning(err)
			break
		}
		glog.Info("--------------------New request----------------------")
		glog.Info("Request: ", string(text))
		req := new(msgs.ClientRequest)
		err = msgs.Unmarshal(text, req)
		if err != nil {
			glog.Fatal(err)
		}

		// construct reply
		reply := handleRequest(*req)
		b, err := msgs.Marshal(reply)
		if err != nil {
			glog.Fatal("error:", err)
		}
		glog.Info(string(b))

		// send reply
		// TODO: FIX currently all server send back replies
		glog.Info("Sending ", string(b))
		n, err := writer.Write(b)
		if err != nil {
			glog.Fatal(err)
		}
		_, err = writer.Write([]byte("\n"))
		if err != nil {
			glog.Fatal(err)
		}

		// tidy up
		err = writer.Flush()
		glog.Info("Finished sending ", n, " bytes")

	}

	cn.Close()
}

func main() {
	// set up logging
	flag.Parse()
	defer glog.Flush()

	conf := config.ParseServerConfig(*config_file)
	if *id == -1 {
		glog.Fatal("ID is required")
	}

	glog.Info("Starting server ", *id)
	defer glog.Warning("Shutting down server ", *id)

	//set up state machine
	keyval = store.New()
	c = cache.Create()
	// setup IO
	cons_io = msgs.MakeIo(2000, len(conf.Peers.Address))

	notifyclient = make(map[msgs.ClientRequest](chan msgs.ClientResponse))
	notifyclient_mutex = sync.RWMutex{}
	go stateMachine()

	// setting up persistent log
	disk, disk_reader, disk_file, is_empty := openFile(*disk_path + "/persistent_log_" + strconv.Itoa(*id) + ".temp")
	defer disk.Flush()
	meta_disk, meta_disk_reader, meta_disk_file, is_new := openFile(*disk_path + "/persistent_data_" + strconv.Itoa(*id) + ".temp")
	defer meta_disk.Flush()

	// check persistent storage for commands
	found := false
	log := make([]msgs.Entry, 10000) // TODO: Remove hard coded limit
	log_length := 0

	if !is_empty {
		for {
			b, err := disk_reader.ReadBytes(byte('\n'))
			if err != nil {
				glog.Info("No more commands in persistent storage")
				break
			}
			found = true
			var update msgs.LogUpdate
			err = msgs.Unmarshal(b, &update)
			if err != nil {
				glog.Fatal("Cannot parse log update", err)
			}
			log[update.Index] = update.Entry
			if log_length < update.Index {
				log_length = update.Index
			}
			glog.Info("Adding from persistent storage :", update)
		}
	}
	log = log[:log_length]

	// check persistent storage for view
	view := 0
	if !is_new {
		for {
			b, err := meta_disk_reader.ReadBytes(byte('\n'))
			if err != nil {
				glog.Info("No more view updates in persistent storage")
				break
			}
			found = true
			view, _ = strconv.Atoi(string(b))
		}
	}

	// write updates to persistent storage
	go func() {
		for {
			view := <-cons_io.ViewPersist
			glog.Info("Updating view to ", view)
			_, err := meta_disk.Write([]byte(strconv.Itoa(view)))
			_, err = meta_disk.Write([]byte("\n"))
			if err != nil {
				glog.Fatal(err)
			}
			meta_disk_file.Sync()
			cons_io.ViewPersistFsync <- view
		}
	}()

	go func() {
		for {
			log := <-cons_io.LogPersist
			glog.Info("Updating log with ", log)
			b, err := msgs.Marshal(log)
			if err != nil {
				glog.Fatal(err)
			}
			// write to persistent storage
			n1, err := disk.Write(b)
			n2, err := disk.Write([]byte("\n"))
			if err != nil {
				glog.Fatal(err)
			}
			glog.Info(n1+n2, " bytes written to persistent log")
			disk_file.Sync()
			cons_io.LogPersistFsync <- log
		}
	}()

	// set up client server
	glog.Info("Starting up client server")
	client_port := strings.Split(conf.Clients.Address[*id],":")[1]
	listeningPort := ":" + client_port
	ln, err := net.Listen("tcp", listeningPort)
	if err != nil {
		glog.Fatal(err)
	}

	// handle for incoming clients
	go func() {
		for {
			conn, err := ln.Accept()
			if err != nil {
				glog.Fatal(err)
			}
			go handleConnection(conn)
		}
	}()

	//set up peer state
	peers = make([]Peer, len(conf.Peers.Address))
	for i := range conf.Peers.Address {
		peers[i] = Peer{
			i, conf.Peers.Address[i], false}
	}
	peers_mutex = sync.RWMutex{}

	//set up peer server
	glog.Info("Starting up peer server")
	peer_port := strings.Split(conf.Peers.Address[*id],":")[1]
	listeningPort = ":" + peer_port
	lnPeers, err := net.Listen("tcp", listeningPort)
	if err != nil {
		glog.Fatal(err)
	}

	// handle local peer (without sending network traffic)
	peers_mutex.Lock()
	peers[*id].handled = true
	peers_mutex.Unlock()
	from := &(cons_io.Incoming)
	go from.Forward(cons_io.OutgoingUnicast[*id])

	// handle for incoming peers
	go func() {
		for {
			conn, err := lnPeers.Accept()
			if err != nil {
				glog.Fatal(err)
			}
			go handlePeer(conn, false)
		}
	}()

	// regularly check if all peers are connected and retry if not
	go func() {
		for {
			checkPeer()
			time.Sleep(500 * time.Millisecond)
		}
	}()

	// setting up the consensus algorithm
	log_max_length := 1000
	if conf.Options.Length > 0 {
		log_max_length = conf.Options.Length
	}
	cons_config := consensus.Config{*id, len(conf.Peers.Address),
		log_max_length, conf.Options.BatchInterval, conf.Options.MaxBatch, conf.Options.DelegateReplication, conf.Options.WindowSize}
	if !found {
		glog.Info("Starting fresh consensus instance")
		go consensus.Init(cons_io, cons_config)
	} else {
		glog.Info("Restoring consensus instance")
		go consensus.Recover(cons_io, cons_config, view, log)
	}
	//go cons_io.DumpPersistentStorage()

	// tidy up
	glog.Info("Setup complete")

	// waiting for exit
	// always flush (whatever happens)
	sigs := make(chan os.Signal, 1)
	signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM)
	sig := <-sigs
	disk.Flush()
	meta_disk.Flush()
	glog.Flush()
	glog.Warning("Shutting down due to ", sig)
}
<|endoftext|>


<commit_before>// Copyright 2015 realglobe, Inc.
//
// 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 account

import (
	"github.com/realglobe-Inc/edo-id-provider/database/account"
	"github.com/realglobe-Inc/edo-id-provider/database/pairwise"
	"github.com/realglobe-Inc/edo-id-provider/database/sector"
	"github.com/realglobe-Inc/edo-id-provider/database/token"
	"github.com/realglobe-Inc/edo-id-provider/idputil"
	"github.com/realglobe-Inc/edo-id-provider/scope"
	tadb "github.com/realglobe-Inc/edo-idp-selector/database/ta"
	idperr "github.com/realglobe-Inc/edo-idp-selector/error"
	requtil "github.com/realglobe-Inc/edo-idp-selector/request"
	logutil "github.com/realglobe-Inc/edo-lib/log"
	"github.com/realglobe-Inc/edo-lib/rand"
	"github.com/realglobe-Inc/edo-lib/server"
	"github.com/realglobe-Inc/go-lib/erro"
	"github.com/realglobe-Inc/go-lib/rglog/level"
	"net/http"
)

type handler struct {
	stopper *server.Stopper

	pwSaltLen int

	acntDb account.Db
	taDb   tadb.Db
	sectDb sector.Db
	pwDb   pairwise.Db
	tokDb  token.Db
	idGen  rand.Generator

	debug bool
}

func New(
	stopper *server.Stopper,
	pwSaltLen int,
	acntDb account.Db,
	taDb tadb.Db,
	sectDb sector.Db,
	pwDb pairwise.Db,
	tokDb token.Db,
	idGen rand.Generator,
	debug bool,
) http.Handler {
	return &handler{
		stopper:   stopper,
		pwSaltLen: pwSaltLen,
		acntDb:    acntDb,
		taDb:      taDb,
		sectDb:    sectDb,
		pwDb:      pwDb,
		tokDb:     tokDb,
		idGen:     idGen,
		debug:     debug,
	}
}

func (this *handler) PairwiseSaltLength() int     { return this.pwSaltLen }
func (this *handler) SectorDb() sector.Db         { return this.sectDb }
func (this *handler) PairwiseDb() pairwise.Db     { return this.pwDb }
func (this *handler) IdGenerator() rand.Generator { return this.idGen }

func (this *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	var sender *requtil.Request

	// panic 対策。
	defer func() {
		if rcv := recover(); rcv != nil {
			idperr.RespondJson(w, r, erro.New(rcv), sender)
			return
		}
	}()

	if this.stopper != nil {
		this.stopper.Stop()
		defer this.stopper.Unstop()
	}

	//////////////////////////////
	server.LogRequest(level.DEBUG, r, this.debug)
	//////////////////////////////

	sender = requtil.Parse(r, "")
	log.Info(sender, ": Received account request")
	defer log.Info(sender, ": Handled account request")

	if err := this.serve(w, r, sender); err != nil {
		idperr.RespondJson(w, r, erro.Wrap(err), sender)
		return
	}
}

func (this *handler) serve(w http.ResponseWriter, r *http.Request, sender *requtil.Request) error {
	req, err := parseRequest(r)
	if err != nil {
		return erro.Wrap(idperr.New(idperr.Invalid_request, erro.Unwrap(err).Error(), http.StatusBadRequest, err))
	} else if req.scheme() != tagBearer {
		return erro.Wrap(idperr.New(idperr.Invalid_request, "unsupported authorization scheme "+req.scheme(), http.StatusBadRequest, nil))
	}

	log.Debug(sender, ": Authrization scheme "+req.scheme()+" is OK")

	if req.token() == "" {
		return erro.Wrap(idperr.New(idperr.Invalid_request, "no access token", http.StatusBadRequest, nil))
	}

	log.Debug(sender, ": Access token "+logutil.Mosaic(req.token())+" is declared")

	tok, err := this.tokDb.Get(req.token())
	if err != nil {
		return erro.Wrap(err)
	} else if tok == nil {
		return erro.Wrap(idperr.New(idperr.Invalid_token, "token is not exist", http.StatusBadRequest, nil))
	} else if tok.Invalid() {
		return erro.Wrap(idperr.New(idperr.Invalid_token, "token is invalid", http.StatusBadRequest, nil))
	}

	log.Debug(sender, ": Declared access token "+logutil.Mosaic(tok.Id())+" is OK")

	ta, err := this.taDb.Get(tok.Ta())
	if err != nil {
		return erro.Wrap(err)
	} else if ta == nil {
		return erro.Wrap(idperr.New(idperr.Invalid_token, "TA "+tok.Ta()+" is not exist", http.StatusBadRequest, nil))
	}

	log.Debug(sender, ": TA "+ta.Id()+" is exist")

	acnt, err := this.acntDb.Get(tok.Account())
	if err != nil {
		return erro.Wrap(err)
	} else if acnt == nil {
		return erro.Wrap(idperr.New(idperr.Invalid_token, "account is not exist", http.StatusBadRequest, nil))
	}

	log.Debug(sender, ": Account "+acnt.Id()+" is exist")

	if err := idputil.SetSub(this, acnt, ta); err != nil {
		return erro.Wrap(err)
	}

	attrNames := map[string]bool{}
	for scop := range tok.Scope() {
		for attrName := range scope.Attributes(scop) {
			attrNames[attrName] = true
		}
	}
	for attrName := range tok.Attributes() {
		attrNames[attrName] = true
	}
	attrNames[tagSub] = true

	log.Debug(sender, ": Return attributes ", attrNames)

	attrs := map[string]interface{}{}
	for attrName := range attrNames {
		attr := acnt.Attribute(attrName)
		if attr == nil || attr == "" {
			continue
		}
		attrs[attrName] = attr
	}

	return idputil.RespondJson(w, attrs)
}
<commit_msg>アカウント情報エンドポイントではスコープから属性を引っ張ってこないようにした<commit_after>// Copyright 2015 realglobe, Inc.
//
// 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 account

import (
	"github.com/realglobe-Inc/edo-id-provider/database/account"
	"github.com/realglobe-Inc/edo-id-provider/database/pairwise"
	"github.com/realglobe-Inc/edo-id-provider/database/sector"
	"github.com/realglobe-Inc/edo-id-provider/database/token"
	"github.com/realglobe-Inc/edo-id-provider/idputil"
	tadb "github.com/realglobe-Inc/edo-idp-selector/database/ta"
	idperr "github.com/realglobe-Inc/edo-idp-selector/error"
	requtil "github.com/realglobe-Inc/edo-idp-selector/request"
	logutil "github.com/realglobe-Inc/edo-lib/log"
	"github.com/realglobe-Inc/edo-lib/rand"
	"github.com/realglobe-Inc/edo-lib/server"
	"github.com/realglobe-Inc/go-lib/erro"
	"github.com/realglobe-Inc/go-lib/rglog/level"
	"net/http"
)

type handler struct {
	stopper *server.Stopper

	pwSaltLen int

	acntDb account.Db
	taDb   tadb.Db
	sectDb sector.Db
	pwDb   pairwise.Db
	tokDb  token.Db
	idGen  rand.Generator

	debug bool
}

func New(
	stopper *server.Stopper,
	pwSaltLen int,
	acntDb account.Db,
	taDb tadb.Db,
	sectDb sector.Db,
	pwDb pairwise.Db,
	tokDb token.Db,
	idGen rand.Generator,
	debug bool,
) http.Handler {
	return &handler{
		stopper:   stopper,
		pwSaltLen: pwSaltLen,
		acntDb:    acntDb,
		taDb:      taDb,
		sectDb:    sectDb,
		pwDb:      pwDb,
		tokDb:     tokDb,
		idGen:     idGen,
		debug:     debug,
	}
}

func (this *handler) PairwiseSaltLength() int     { return this.pwSaltLen }
func (this *handler) SectorDb() sector.Db         { return this.sectDb }
func (this *handler) PairwiseDb() pairwise.Db     { return this.pwDb }
func (this *handler) IdGenerator() rand.Generator { return this.idGen }

func (this *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	var sender *requtil.Request

	// panic 対策。
	defer func() {
		if rcv := recover(); rcv != nil {
			idperr.RespondJson(w, r, erro.New(rcv), sender)
			return
		}
	}()

	if this.stopper != nil {
		this.stopper.Stop()
		defer this.stopper.Unstop()
	}

	//////////////////////////////
	server.LogRequest(level.DEBUG, r, this.debug)
	//////////////////////////////

	sender = requtil.Parse(r, "")
	log.Info(sender, ": Received account request")
	defer log.Info(sender, ": Handled account request")

	if err := this.serve(w, r, sender); err != nil {
		idperr.RespondJson(w, r, erro.Wrap(err), sender)
		return
	}
}

func (this *handler) serve(w http.ResponseWriter, r *http.Request, sender *requtil.Request) error {
	req, err := parseRequest(r)
	if err != nil {
		return erro.Wrap(idperr.New(idperr.Invalid_request, erro.Unwrap(err).Error(), http.StatusBadRequest, err))
	} else if req.scheme() != tagBearer {
		return erro.Wrap(idperr.New(idperr.Invalid_request, "unsupported authorization scheme "+req.scheme(), http.StatusBadRequest, nil))
	}

	log.Debug(sender, ": Authrization scheme "+req.scheme()+" is OK")

	if req.token() == "" {
		return erro.Wrap(idperr.New(idperr.Invalid_request, "no access token", http.StatusBadRequest, nil))
	}

	log.Debug(sender, ": Access token "+logutil.Mosaic(req.token())+" is declared")

	tok, err := this.tokDb.Get(req.token())
	if err != nil {
		return erro.Wrap(err)
	} else if tok == nil {
		return erro.Wrap(idperr.New(idperr.Invalid_token, "token is not exist", http.StatusBadRequest, nil))
	} else if tok.Invalid() {
		return erro.Wrap(idperr.New(idperr.Invalid_token, "token is invalid", http.StatusBadRequest, nil))
	}

	log.Debug(sender, ": Declared access token "+logutil.Mosaic(tok.Id())+" is OK")

	ta, err := this.taDb.Get(tok.Ta())
	if err != nil {
		return erro.Wrap(err)
	} else if ta == nil {
		return erro.Wrap(idperr.New(idperr.Invalid_token, "TA "+tok.Ta()+" is not exist", http.StatusBadRequest, nil))
	}

	log.Debug(sender, ": TA "+ta.Id()+" is exist")

	acnt, err := this.acntDb.Get(tok.Account())
	if err != nil {
		return erro.Wrap(err)
	} else if acnt == nil {
		return erro.Wrap(idperr.New(idperr.Invalid_token, "account is not exist", http.StatusBadRequest, nil))
	}

	log.Debug(sender, ": Account "+acnt.Id()+" is exist")

	if err := idputil.SetSub(this, acnt, ta); err != nil {
		return erro.Wrap(err)
	}

	attrNames := map[string]bool{}
	for attrName := range tok.Attributes() {
		attrNames[attrName] = true
	}
	attrNames[tagSub] = true

	log.Debug(sender, ": Return attributes ", attrNames)

	attrs := map[string]interface{}{}
	for attrName := range attrNames {
		attr := acnt.Attribute(attrName)
		if attr == nil || attr == "" {
			continue
		}
		attrs[attrName] = attr
	}

	return idputil.RespondJson(w, attrs)
}
<|endoftext|>


<commit_before>/*
Copyright 2016 The Kubernetes Authors.

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.
*/

// Code generated by ""fitask" -type=IAMInstanceProfileRole"; DO NOT EDIT

package awstasks

import (
	"encoding/json"

	"k8s.io/kops/upup/pkg/fi"
)

// IAMInstanceProfileRole

// JSON marshalling boilerplate
type realIAMInstanceProfileRole IAMInstanceProfileRole

// UnmarshalJSON implements conversion to JSON, supporitng an alternate specification of the object as a string
func (o *IAMInstanceProfileRole) UnmarshalJSON(data []byte) error {
	var jsonName string
	if err := json.Unmarshal(data, &jsonName); err == nil {
		o.Name = &jsonName
		return nil
	}

	var r realIAMInstanceProfileRole
	if err := json.Unmarshal(data, &r); err != nil {
		return err
	}
	*o = IAMInstanceProfileRole(r)
	return nil
}

var _ fi.HasName = &IAMInstanceProfileRole{}

// GetName returns the Name of the object, implementing fi.HasName
func (o *IAMInstanceProfileRole) GetName() *string {
	return o.Name
}

// SetName sets the Name of the object, implementing fi.SetName
func (o *IAMInstanceProfileRole) SetName(name string) {
	o.Name = &name
}

// String is the stringer function for the task, producing readable output using fi.TaskAsString
func (o *IAMInstanceProfileRole) String() string {
	return fi.TaskAsString(o)
}
<commit_msg>Simplify terraform ELB tasks<commit_after><|endoftext|>


<commit_before>// Copyright 2015/2016 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package prog

import (
	"fmt"
	"math/rand"
	"sort"
)

// Calulation of call-to-call priorities.
// For a given pair of calls X and Y, the priority is our guess as to whether
// additional of call Y into a program containing call X is likely to give
// new coverage or not.
// The current algorithm has two components: static and dynamic.
// The static component is based on analysis of argument types. For example,
// if call X and call Y both accept fd[sock], then they are more likely to give
// new coverage together.
// The dynamic component is based on frequency of occurrence of a particular
// pair of syscalls in a single program in corpus. For example, if socket and
// connect frequently occur in programs together, we give higher priority to
// this pair of syscalls.
// Note: the current implementation is very basic, there is no theory behind any
// constants.

func (target *Target) CalculatePriorities(corpus []*Prog) [][]float32 {
	static := target.calcStaticPriorities()
	dynamic := target.calcDynamicPrio(corpus)
	for i, prios := range static {
		for j, p := range prios {
			dynamic[i][j] *= p
		}
	}
	return dynamic
}

func (target *Target) calcStaticPriorities() [][]float32 {
	uses := make(map[string]map[int]float32)
	for _, c := range target.Syscalls {
		noteUsage := func(weight float32, str string, args ...interface{}) {
			id := fmt.Sprintf(str, args...)
			if uses[id] == nil {
				uses[id] = make(map[int]float32)
			}
			old := uses[id][c.ID]
			if weight > old {
				uses[id][c.ID] = weight
			}
		}
		ForeachType(c, func(t Type) {
			switch a := t.(type) {
			case *ResourceType:
				if a.Desc.Name == "pid" || a.Desc.Name == "uid" || a.Desc.Name == "gid" {
					// Pid/uid/gid usually play auxiliary role,
					// but massively happen in some structs.
					noteUsage(0.1, "res%v", a.Desc.Name)
				} else {
					str := "res"
					for i, k := range a.Desc.Kind {
						str += "-" + k
						w := 1.0
						if i < len(a.Desc.Kind)-1 {
							w = 0.2
						}
						noteUsage(float32(w), str)
					}
				}
			case *PtrType:
				if _, ok := a.Type.(*StructType); ok {
					noteUsage(1.0, "ptrto-%v", a.Type.Name())
				}
				if _, ok := a.Type.(*UnionType); ok {
					noteUsage(1.0, "ptrto-%v", a.Type.Name())
				}
				if arr, ok := a.Type.(*ArrayType); ok {
					noteUsage(1.0, "ptrto-%v", arr.Type.Name())
				}
			case *BufferType:
				switch a.Kind {
				case BufferBlobRand, BufferBlobRange, BufferText:
				case BufferString:
					if a.SubKind != "" {
						noteUsage(0.2, fmt.Sprintf("str-%v", a.SubKind))
					}
				case BufferFilename:
					noteUsage(1.0, "filename")
				default:
					panic("unknown buffer kind")
				}
			case *VmaType:
				noteUsage(0.5, "vma")
			case *IntType:
				switch a.Kind {
				case IntPlain, IntFileoff, IntRange:
				default:
					panic("unknown int kind")
				}
			}
		})
	}
	prios := make([][]float32, len(target.Syscalls))
	for i := range prios {
		prios[i] = make([]float32, len(target.Syscalls))
	}
	for _, calls := range uses {
		for c0, w0 := range calls {
			for c1, w1 := range calls {
				if c0 == c1 {
					// Self-priority is assigned below.
					continue
				}
				prios[c0][c1] += w0 * w1
			}
		}
	}

	// Self-priority (call wrt itself) is assigned to the maximum priority
	// this call has wrt other calls. This way the priority is high, but not too high.
	for c0, pp := range prios {
		var max float32
		for _, p := range pp {
			if max < p {
				max = p
			}
		}
		pp[c0] = max
	}
	normalizePrio(prios)
	return prios
}

func (target *Target) calcDynamicPrio(corpus []*Prog) [][]float32 {
	prios := make([][]float32, len(target.Syscalls))
	for i := range prios {
		prios[i] = make([]float32, len(target.Syscalls))
	}
	for _, p := range corpus {
		for _, c0 := range p.Calls {
			for _, c1 := range p.Calls {
				id0 := c0.Meta.ID
				id1 := c1.Meta.ID
				prios[id0][id1] += 1.0
			}
		}
	}
	normalizePrio(prios)
	return prios
}

// normalizePrio assigns some minimal priorities to calls with zero priority,
// and then normalizes priorities to 0.1..1 range.
func normalizePrio(prios [][]float32) {
	for _, prio := range prios {
		max := float32(0)
		min := float32(1e10)
		nzero := 0
		for _, p := range prio {
			if max < p {
				max = p
			}
			if p != 0 && min > p {
				min = p
			}
			if p == 0 {
				nzero++
			}
		}
		if nzero != 0 {
			min /= 2 * float32(nzero)
		}
		for i, p := range prio {
			if max == 0 {
				prio[i] = 1
				continue
			}
			if p == 0 {
				p = min
			}
			p = (p-min)/(max-min)*0.9 + 0.1
			if p > 1 {
				p = 1
			}
			prio[i] = p
		}
	}
}

// ChooseTable allows to do a weighted choice of a syscall for a given syscall
// based on call-to-call priorities and a set of enabled syscalls.
type ChoiceTable struct {
	target       *Target
	run          [][]int
	enabledCalls []*Syscall
	enabled      map[*Syscall]bool
}

func (target *Target) BuildChoiceTable(prios [][]float32, enabled map[*Syscall]bool) *ChoiceTable {
	if enabled == nil {
		enabled = make(map[*Syscall]bool)
		for _, c := range target.Syscalls {
			enabled[c] = true
		}
	}
	var enabledCalls []*Syscall
	for c := range enabled {
		enabledCalls = append(enabledCalls, c)
	}
	run := make([][]int, len(target.Syscalls))
	for i := range run {
		if !enabled[target.Syscalls[i]] {
			continue
		}
		run[i] = make([]int, len(target.Syscalls))
		sum := 0
		for j := range run[i] {
			if enabled[target.Syscalls[j]] {
				w := 1
				if prios != nil {
					w = int(prios[i][j] * 1000)
				}
				sum += w
			}
			run[i][j] = sum
		}
	}
	return &ChoiceTable{target, run, enabledCalls, enabled}
}

func (ct *ChoiceTable) Choose(r *rand.Rand, call int) int {
	if call < 0 {
		return ct.enabledCalls[r.Intn(len(ct.enabledCalls))].ID
	}
	run := ct.run[call]
	if run == nil {
		return ct.enabledCalls[r.Intn(len(ct.enabledCalls))].ID
	}
	for {
		x := r.Intn(run[len(run)-1]) + 1
		i := sort.SearchInts(run, x)
		if ct.enabled[ct.target.Syscalls[i]] {
			return i
		}
	}
}
<commit_msg>prog: refactor calcStaticPriorities<commit_after>// Copyright 2015/2016 syzkaller project authors. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package prog

import (
	"fmt"
	"math/rand"
	"sort"
)

// Calulation of call-to-call priorities.
// For a given pair of calls X and Y, the priority is our guess as to whether
// additional of call Y into a program containing call X is likely to give
// new coverage or not.
// The current algorithm has two components: static and dynamic.
// The static component is based on analysis of argument types. For example,
// if call X and call Y both accept fd[sock], then they are more likely to give
// new coverage together.
// The dynamic component is based on frequency of occurrence of a particular
// pair of syscalls in a single program in corpus. For example, if socket and
// connect frequently occur in programs together, we give higher priority to
// this pair of syscalls.
// Note: the current implementation is very basic, there is no theory behind any
// constants.

func (target *Target) CalculatePriorities(corpus []*Prog) [][]float32 {
	static := target.calcStaticPriorities()
	dynamic := target.calcDynamicPrio(corpus)
	for i, prios := range static {
		for j, p := range prios {
			dynamic[i][j] *= p
		}
	}
	return dynamic
}

func (target *Target) calcStaticPriorities() [][]float32 {
	uses := target.calcResourceUsage()
	prios := make([][]float32, len(target.Syscalls))
	for i := range prios {
		prios[i] = make([]float32, len(target.Syscalls))
	}
	for _, calls := range uses {
		for c0, w0 := range calls {
			for c1, w1 := range calls {
				if c0 == c1 {
					// Self-priority is assigned below.
					continue
				}
				prios[c0][c1] += w0 * w1
			}
		}
	}

	// Self-priority (call wrt itself) is assigned to the maximum priority
	// this call has wrt other calls. This way the priority is high, but not too high.
	for c0, pp := range prios {
		var max float32
		for _, p := range pp {
			if max < p {
				max = p
			}
		}
		pp[c0] = max
	}
	normalizePrio(prios)
	return prios
}

func (target *Target) calcResourceUsage() map[string]map[int]float32 {
	uses := make(map[string]map[int]float32)
	for _, c := range target.Syscalls {
		ForeachType(c, func(t Type) {
			switch a := t.(type) {
			case *ResourceType:
				if a.Desc.Name == "pid" || a.Desc.Name == "uid" || a.Desc.Name == "gid" {
					// Pid/uid/gid usually play auxiliary role,
					// but massively happen in some structs.
					noteUsage(uses, c, 0.1, "res%v", a.Desc.Name)
				} else {
					str := "res"
					for i, k := range a.Desc.Kind {
						str += "-" + k
						w := 1.0
						if i < len(a.Desc.Kind)-1 {
							w = 0.2
						}
						noteUsage(uses, c, float32(w), str)
					}
				}
			case *PtrType:
				if _, ok := a.Type.(*StructType); ok {
					noteUsage(uses, c, 1.0, "ptrto-%v", a.Type.Name())
				}
				if _, ok := a.Type.(*UnionType); ok {
					noteUsage(uses, c, 1.0, "ptrto-%v", a.Type.Name())
				}
				if arr, ok := a.Type.(*ArrayType); ok {
					noteUsage(uses, c, 1.0, "ptrto-%v", arr.Type.Name())
				}
			case *BufferType:
				switch a.Kind {
				case BufferBlobRand, BufferBlobRange, BufferText:
				case BufferString:
					if a.SubKind != "" {
						noteUsage(uses, c, 0.2, fmt.Sprintf("str-%v", a.SubKind))
					}
				case BufferFilename:
					noteUsage(uses, c, 1.0, "filename")
				default:
					panic("unknown buffer kind")
				}
			case *VmaType:
				noteUsage(uses, c, 0.5, "vma")
			case *IntType:
				switch a.Kind {
				case IntPlain, IntFileoff, IntRange:
				default:
					panic("unknown int kind")
				}
			}
		})
	}
	return uses
}

func noteUsage(uses map[string]map[int]float32, c *Syscall, weight float32, str string, args ...interface{}) {
	id := fmt.Sprintf(str, args...)
	if uses[id] == nil {
		uses[id] = make(map[int]float32)
	}
	old := uses[id][c.ID]
	if weight > old {
		uses[id][c.ID] = weight
	}
}

func (target *Target) calcDynamicPrio(corpus []*Prog) [][]float32 {
	prios := make([][]float32, len(target.Syscalls))
	for i := range prios {
		prios[i] = make([]float32, len(target.Syscalls))
	}
	for _, p := range corpus {
		for _, c0 := range p.Calls {
			for _, c1 := range p.Calls {
				id0 := c0.Meta.ID
				id1 := c1.Meta.ID
				prios[id0][id1] += 1.0
			}
		}
	}
	normalizePrio(prios)
	return prios
}

// normalizePrio assigns some minimal priorities to calls with zero priority,
// and then normalizes priorities to 0.1..1 range.
func normalizePrio(prios [][]float32) {
	for _, prio := range prios {
		max := float32(0)
		min := float32(1e10)
		nzero := 0
		for _, p := range prio {
			if max < p {
				max = p
			}
			if p != 0 && min > p {
				min = p
			}
			if p == 0 {
				nzero++
			}
		}
		if nzero != 0 {
			min /= 2 * float32(nzero)
		}
		for i, p := range prio {
			if max == 0 {
				prio[i] = 1
				continue
			}
			if p == 0 {
				p = min
			}
			p = (p-min)/(max-min)*0.9 + 0.1
			if p > 1 {
				p = 1
			}
			prio[i] = p
		}
	}
}

// ChooseTable allows to do a weighted choice of a syscall for a given syscall
// based on call-to-call priorities and a set of enabled syscalls.
type ChoiceTable struct {
	target       *Target
	run          [][]int
	enabledCalls []*Syscall
	enabled      map[*Syscall]bool
}

func (target *Target) BuildChoiceTable(prios [][]float32, enabled map[*Syscall]bool) *ChoiceTable {
	if enabled == nil {
		enabled = make(map[*Syscall]bool)
		for _, c := range target.Syscalls {
			enabled[c] = true
		}
	}
	var enabledCalls []*Syscall
	for c := range enabled {
		enabledCalls = append(enabledCalls, c)
	}
	run := make([][]int, len(target.Syscalls))
	for i := range run {
		if !enabled[target.Syscalls[i]] {
			continue
		}
		run[i] = make([]int, len(target.Syscalls))
		sum := 0
		for j := range run[i] {
			if enabled[target.Syscalls[j]] {
				w := 1
				if prios != nil {
					w = int(prios[i][j] * 1000)
				}
				sum += w
			}
			run[i][j] = sum
		}
	}
	return &ChoiceTable{target, run, enabledCalls, enabled}
}

func (ct *ChoiceTable) Choose(r *rand.Rand, call int) int {
	if call < 0 {
		return ct.enabledCalls[r.Intn(len(ct.enabledCalls))].ID
	}
	run := ct.run[call]
	if run == nil {
		return ct.enabledCalls[r.Intn(len(ct.enabledCalls))].ID
	}
	for {
		x := r.Intn(run[len(run)-1]) + 1
		i := sort.SearchInts(run, x)
		if ct.enabled[ct.target.Syscalls[i]] {
			return i
		}
	}
}
<|endoftext|>


<commit_before>package version

// Version of Functions
var Version = "0.3.200"
<commit_msg>fnserver: 0.3.201 release [skip ci]<commit_after>package version

// Version of Functions
var Version = "0.3.201"
<|endoftext|>


<commit_before>package version

// Version of Functions
var Version = "0.3.620"
<commit_msg>fnserver: 0.3.621 release [skip ci]<commit_after>package version

// Version of Functions
var Version = "0.3.621"
<|endoftext|>


<commit_before>package version

// Version of Functions
var Version = "0.3.329"
<commit_msg>fnserver: 0.3.330 release [skip ci]<commit_after>package version

// Version of Functions
var Version = "0.3.330"
<|endoftext|>


<commit_before>package version

// Version of Functions
var Version = "0.3.293"
<commit_msg>fnserver: 0.3.294 release [skip ci]<commit_after>package version

// Version of Functions
var Version = "0.3.294"
<|endoftext|>


<commit_before>package version

// Version of Functions
var Version = "0.3.143"
<commit_msg>functions: 0.3.144 release [skip ci]<commit_after>package version

// Version of Functions
var Version = "0.3.144"
<|endoftext|>


<commit_before>package suggest

import (
	"database/sql"
	"time"

	"github.com/remeh/smartwitter/account"
	"github.com/remeh/smartwitter/storage"
	"github.com/remeh/smartwitter/twitter"

	"github.com/lib/pq"
)

func SuggestByKeywords(user *account.User, keywords []string, duration time.Duration, limit int) (twitter.Tweets, twitter.TweetDoneActions, error) {
	var rows *sql.Rows
	var err error

	// get tweets with these keywords
	// ----------------------

	ct := time.Now().Add(-duration)

	if rows, err = storage.DB().Query(`
		select
			"tweet"."text",
			"tweet"."twitter_id",
			"tweet"."retweet_count",
			"tweet"."favorite_count",
			"tweet"."link",
			"tweet"."twitter_user_uid",
			"tweet_done_action"."ignored_time",
			"tweet_done_action"."liked_time",
			"tweet_done_action"."retweeted_time",
			"tweet_done_action"."ignored_time" IS NOT NULL,
			"tweet_done_action"."liked_time" IS NOT NULL,
			"tweet_done_action"."retweeted_time" IS NOT NULL
		from "tweet"
		join "twitter_user" on
			"tweet"."twitter_user_uid" = "twitter_user"."uid"
		left join "tweet_done_action" on
			"tweet_done_action"."user_uid" = $3
			and
			"tweet"."twitter_id" = "tweet_done_action"."tweet_id"
		where
			not "text" LIKE 'RT @%'
			and
			"tweet"."keywords" @> $1::text[]
			and
			"tweet"."creation_time" > $2
			and
			"tweet_done_action"."ignored_time" IS NULL
		order by (favorite_count+retweet_count+followers_count) desc, "tweet".uid
		limit $4;
	`, pq.Array(keywords), ct, user.Uid, limit); err != nil {
		return nil, nil, err
	}

	defer rows.Close()

	rv := make(twitter.Tweets, 0)
	tdas := make(twitter.TweetDoneActions, 0)

	for rows.Next() {
		t := twitter.Tweet{}
		tda := twitter.TweetDoneAction{}

		var it, lt, rt pq.NullTime

		if err := rows.Scan(
			&t.Text,
			&t.TwitterId,
			&t.RetweetCount,
			&t.FavoriteCount,
			&t.Link,
			&t.TwitterUserUid,
			&it,
			&lt,
			&rt,
			&tda.Ignored,
			&tda.Liked,
			&tda.Retweeted,
		); err != nil {
			return nil, nil, err
		}

		if it.Valid {
			tda.IgnoreTime = it.Time
		}
		if lt.Valid {
			tda.LikeTime = lt.Time
		}
		if rt.Valid {
			tda.RetweetTime = rt.Time
		}

		rv = append(rv, t)
		tdas = append(tdas, tda)
	}

	return rv, tdas, nil
}
<commit_msg>suggest: remove the followers count in the suggestion algo.<commit_after>package suggest

import (
	"database/sql"
	"time"

	"github.com/remeh/smartwitter/account"
	"github.com/remeh/smartwitter/storage"
	"github.com/remeh/smartwitter/twitter"

	"github.com/lib/pq"
)

func SuggestByKeywords(user *account.User, keywords []string, duration time.Duration, limit int) (twitter.Tweets, twitter.TweetDoneActions, error) {
	var rows *sql.Rows
	var err error

	// get tweets with these keywords
	// ----------------------

	ct := time.Now().Add(-duration)

	if rows, err = storage.DB().Query(`
		select
			"tweet"."text",
			"tweet"."twitter_id",
			"tweet"."retweet_count",
			"tweet"."favorite_count",
			"tweet"."link",
			"tweet"."twitter_user_uid",
			"tweet_done_action"."ignored_time",
			"tweet_done_action"."liked_time",
			"tweet_done_action"."retweeted_time",
			"tweet_done_action"."ignored_time" IS NOT NULL,
			"tweet_done_action"."liked_time" IS NOT NULL,
			"tweet_done_action"."retweeted_time" IS NOT NULL
		from "tweet"
		join "twitter_user" on
			"tweet"."twitter_user_uid" = "twitter_user"."uid"
		left join "tweet_done_action" on
			"tweet_done_action"."user_uid" = $3
			and
			"tweet"."twitter_id" = "tweet_done_action"."tweet_id"
		where
			not "text" LIKE 'RT @%'
			and
			"tweet"."keywords" @> $1::text[]
			and
			"tweet"."creation_time" > $2
			and
			"tweet_done_action"."ignored_time" IS NULL
		order by (favorite_count+retweet_count) desc, "tweet".uid
		limit $4;
	`, pq.Array(keywords), ct, user.Uid, limit); err != nil {
		return nil, nil, err
	}

	defer rows.Close()

	rv := make(twitter.Tweets, 0)
	tdas := make(twitter.TweetDoneActions, 0)

	for rows.Next() {
		t := twitter.Tweet{}
		tda := twitter.TweetDoneAction{}

		var it, lt, rt pq.NullTime

		if err := rows.Scan(
			&t.Text,
			&t.TwitterId,
			&t.RetweetCount,
			&t.FavoriteCount,
			&t.Link,
			&t.TwitterUserUid,
			&it,
			&lt,
			&rt,
			&tda.Ignored,
			&tda.Liked,
			&tda.Retweeted,
		); err != nil {
			return nil, nil, err
		}

		if it.Valid {
			tda.IgnoreTime = it.Time
		}
		if lt.Valid {
			tda.LikeTime = lt.Time
		}
		if rt.Valid {
			tda.RetweetTime = rt.Time
		}

		rv = append(rv, t)
		tdas = append(tdas, tda)
	}

	return rv, tdas, nil
}
<|endoftext|>