<commit_before>package main

import (
	"github.com/subgraph/fw-daemon/proc"
	"github.com/yawning/bulb"

	"bufio"
	"fmt"
	"log"
	"net"
	"strconv"
	"strings"
	"sync"
)

// ProxyListener is used to listen for
// Tor Control port connections and
// dispatches them to worker sessions
// to implement the filtering proxy pipeline.
type ProxyListener struct {
	cfg         *RoflcoptorConfig
	watch       bool
	wg          *sync.WaitGroup
	tcpListener *net.TCPListener
	errChan     chan error
}

// NewProxyListener creates a new ProxyListener given
// a configuration structure.
func NewProxyListener(cfg *RoflcoptorConfig, wg *sync.WaitGroup, watch bool) (*ProxyListener, error) {
	p := ProxyListener{
		cfg:   cfg,
		wg:    wg,
		watch: watch,
	}
	return &p, nil
}

// InitAllListeners initialize all tor control port proxy listeners
func (p *ProxyListener) InitAllListeners() {
	// XXX TODO add UNIX domain socket listener
	p.wg.Add(1)
	p.FilterTcpAcceptLoop()
}

// FilterAcceptLoop and listens and accepts new
// connections and passes them into our filter proxy pipeline
func (p *ProxyListener) FilterTcpAcceptLoop() {
	defer p.wg.Done()

	tcpAddr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%s", p.cfg.ListenIP, p.cfg.ListenTCPPort))
	if err != nil {
		panic(fmt.Sprintf("Failed to resolve TCP configured filter port: %s\n", err))
	}
	p.tcpListener, err = net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		panic(fmt.Sprintf("Failed to listen on the filter port: %s\n", err))
	}
	defer p.tcpListener.Close()

	// Listen for incoming connections, and dispatch workers.
	// XXX TODO use a select statement to implement a stop/shutdown channel?
	for {
		conn, err := p.tcpListener.Accept()
		if err != nil {
			if e, ok := err.(net.Error); ok && !e.Temporary() {
				log.Printf("ERR/tor: Failed to Accept(): %v", err)
			}
		}

		// Create the appropriate session instance.
		s := NewProxySession(p.cfg, conn, p.watch)
		go s.sessionWorker()
	}
}

// ProxySession is used to manage is single Tor Control port filter proxy session
type ProxySession struct {
	cfg                *RoflcoptorConfig
	clientFilterPolicy *FilterConfig
	serverFilterPolicy *FilterConfig

	appConn          net.Conn
	appConnReader    *bufio.Reader
	appConnWriteLock sync.Mutex

	torConn   *bulb.Conn
	protoInfo *bulb.ProtocolInfo

	watch   bool
	errChan chan error
	sync.WaitGroup
}

func NewProxySession(cfg *RoflcoptorConfig, conn net.Conn, watch bool) *ProxySession {
	s := &ProxySession{
		cfg:           cfg,
		watch:         watch,
		appConn:       conn,
		appConnReader: bufio.NewReader(conn),
		errChan:       make(chan error, 2),
	}
	return s
}

// getFilterPolicy returns a *ServerClientFilterConfig
// (session policy) if one can be found, otherwise nil is returned.
// Note that the connection is closed upon fatal error,
// when we fail to lookup the proc info, we return nil and an error.
func (s *ProxySession) getFilterPolicy() (*ServerClientFilterConfig, error) {

	fields := strings.Split(s.appConn.RemoteAddr().String(), ":")
	dstPortStr := fields[1]
	dstIP := net.ParseIP(s.cfg.ListenIP)

	if dstIP == nil {
		s.appConn.Close()
		return nil, fmt.Errorf("net.ParseIP fail for: %s\n", s.cfg.ListenIP)
	}

	srcP, _ := strconv.ParseUint(dstPortStr, 10, 16)
	dstP, _ := strconv.ParseUint(s.cfg.ListenTCPPort, 10, 16)
	procInfo := proc.LookupTCPSocketProcess(uint16(srcP), dstIP, uint16(dstP))
	if procInfo == nil {
		s.appConn.Close()
		return nil, fmt.Errorf("Could not find process information for: %d %s %d\n", srcP, dstIP, dstP)
	}

	filter := getFilterForPathAndUID(procInfo.ExePath, procInfo.Uid)
	if filter == nil {
		filter = getFilterForPath(procInfo.ExePath)
	} else {
		log.Printf("No filters found for: %s (%d)\n", procInfo.ExePath, procInfo.Uid)
		filter = nil
	}
	return filter, nil
}

func (s *ProxySession) initTorControl() error {
	var err error

	// Connect to the real control port.
	if s.torConn, err = bulb.Dial("unix", s.cfg.TorControlSocketPath); err != nil {
		return fmt.Errorf("ERR/tor: Failed to connect to tor control port: %v", err)
	}

	// Issue a PROTOCOLINFO, so we can send a realistic response.
	if s.protoInfo, err = s.torConn.ProtocolInfo(); err != nil {
		s.torConn.Close()
		return fmt.Errorf("ERR/tor: Failed to issue PROTOCOLINFO: %v", err)
	}

	// Authenticate with the real tor control port.
	// XXX: Pull password out of `b.s.cfg`.
	if err = s.torConn.Authenticate(""); err != nil {
		s.torConn.Close()
		return fmt.Errorf("ERR/tor: Failed to authenticate: %v", err)
	}
	return nil
}

func (s *ProxySession) sessionWorker() {
	defer s.appConn.Close()

	var policy *ServerClientFilterConfig
	var err error

	clientAddr := s.appConn.RemoteAddr()
	log.Printf("INFO/tor: New ctrl connection from: %s", clientAddr)

	policy, err = s.getFilterPolicy()
	if err != nil {
		log.Printf("proc info query failure; connection from %s aborted: %s\n", clientAddr, err)
		return
	}
	if policy == nil {
		s.clientFilterPolicy = nil
		s.serverFilterPolicy = nil
	} else {
		s.clientFilterPolicy = &FilterConfig{
			Allowed:             policy.ClientAllowed,
			AllowedPrefixes:     policy.ClientAllowedPrefixes,
			Replacements:        policy.ClientReplacements,
			ReplacementPrefixes: policy.ClientReplacementPrefixes,
		}
		s.serverFilterPolicy = &FilterConfig{
			Allowed:             policy.ServerAllowed,
			AllowedPrefixes:     policy.ServerAllowedPrefixes,
			Replacements:        policy.ServerReplacements,
			ReplacementPrefixes: policy.ServerReplacementPrefixes,
		}
	}

	// Authenticate with the real control port
	err = s.initTorControl()
	defer s.torConn.Close()

	if err != nil {
		log.Printf("RoflcopTor: Failed to authenticate with the tor control port: %s\n", err)
		return
	}

	s.appConnReader = bufio.NewReader(s.appConn)
	s.Add(2)
	go s.proxyFilterTorToApp()
	go s.proxyFilterAppToTor()

	// Wait till all sessions are finished, log and return.
	s.Wait()
	s.torConn.Close()
	s.appConn.Close()

	if len(s.errChan) > 0 {
		err := <-s.errChan
		log.Printf("INFO/tor: Closed client connection from: %s: %v", clientAddr, err)
	} else {
		log.Printf("INFO/tor: Closed client connection from: %v", clientAddr)
	}
}

func (s *ProxySession) proxyFilterTorToApp() {
	defer s.Done()

	var appConnLock sync.Mutex
	writeAppConn := func(b []byte) (int, error) {
		appConnLock.Lock() // XXX is the lock really needed here?
		defer appConnLock.Unlock()
		return s.appConn.Write(b)
	}

	rd := bufio.NewReader(s.torConn)
	for {
		line, err := rd.ReadBytes('\n')
		if err != nil {
			s.errChan <- err
			break
		}
		lineStr := strings.TrimSpace(string(line))
		log.Printf("A<-T: [%s]\n", lineStr)
		if s.watch {
			if _, err = writeAppConn([]byte(lineStr + "\n")); err != nil {
				s.errChan <- err
				break
			}
		} else {
			filterCommand(lineStr, "250 OK", writeAppConn, s.errChan, s.serverFilterPolicy)
		}
	}
}

func (s *ProxySession) proxyFilterAppToTor() {
	defer s.Done()

	writeToTor := func(line []byte) (int, error) {
		n, err := s.torConn.Write([]byte(line))
		return n, err
	}

	for {
		line, err := s.appConnReader.ReadBytes('\n')
		if err != nil {
			s.errChan <- err
			break
		}
		lineStr := strings.TrimSpace(string(line))
		log.Printf("A->T: [%s]\n", lineStr)

		if s.watch {
			_, err = writeToTor([]byte(lineStr + "\n"))
			if err != nil {
				s.errChan <- err
			}
		} else {
			filterCommand(lineStr, "", writeToTor, s.errChan, s.clientFilterPolicy)
		}
	}
}
<commit_msg>Print exec path when we receive proc info for a client<commit_after>package main

import (
	"github.com/subgraph/fw-daemon/proc"
	"github.com/yawning/bulb"

	"bufio"
	"fmt"
	"log"
	"net"
	"strconv"
	"strings"
	"sync"
)

// ProxyListener is used to listen for
// Tor Control port connections and
// dispatches them to worker sessions
// to implement the filtering proxy pipeline.
type ProxyListener struct {
	cfg         *RoflcoptorConfig
	watch       bool
	wg          *sync.WaitGroup
	tcpListener *net.TCPListener
	errChan     chan error
}

// NewProxyListener creates a new ProxyListener given
// a configuration structure.
func NewProxyListener(cfg *RoflcoptorConfig, wg *sync.WaitGroup, watch bool) (*ProxyListener, error) {
	p := ProxyListener{
		cfg:   cfg,
		wg:    wg,
		watch: watch,
	}
	return &p, nil
}

// InitAllListeners initialize all tor control port proxy listeners
func (p *ProxyListener) InitAllListeners() {
	// XXX TODO add UNIX domain socket listener
	p.wg.Add(1)
	p.FilterTcpAcceptLoop()
}

// FilterAcceptLoop and listens and accepts new
// connections and passes them into our filter proxy pipeline
func (p *ProxyListener) FilterTcpAcceptLoop() {
	defer p.wg.Done()

	tcpAddr, err := net.ResolveTCPAddr("tcp", fmt.Sprintf("%s:%s", p.cfg.ListenIP, p.cfg.ListenTCPPort))
	if err != nil {
		panic(fmt.Sprintf("Failed to resolve TCP configured filter port: %s\n", err))
	}
	p.tcpListener, err = net.ListenTCP("tcp", tcpAddr)
	if err != nil {
		panic(fmt.Sprintf("Failed to listen on the filter port: %s\n", err))
	}
	defer p.tcpListener.Close()

	// Listen for incoming connections, and dispatch workers.
	// XXX TODO use a select statement to implement a stop/shutdown channel?
	for {
		conn, err := p.tcpListener.Accept()
		if err != nil {
			if e, ok := err.(net.Error); ok && !e.Temporary() {
				log.Printf("ERR/tor: Failed to Accept(): %v", err)
			}
		}

		// Create the appropriate session instance.
		s := NewProxySession(p.cfg, conn, p.watch)
		go s.sessionWorker()
	}
}

// ProxySession is used to manage is single Tor Control port filter proxy session
type ProxySession struct {
	cfg                *RoflcoptorConfig
	clientFilterPolicy *FilterConfig
	serverFilterPolicy *FilterConfig

	appConn          net.Conn
	appConnReader    *bufio.Reader
	appConnWriteLock sync.Mutex

	torConn   *bulb.Conn
	protoInfo *bulb.ProtocolInfo

	watch   bool
	errChan chan error
	sync.WaitGroup
}

func NewProxySession(cfg *RoflcoptorConfig, conn net.Conn, watch bool) *ProxySession {
	s := &ProxySession{
		cfg:           cfg,
		watch:         watch,
		appConn:       conn,
		appConnReader: bufio.NewReader(conn),
		errChan:       make(chan error, 2),
	}
	return s
}

// getFilterPolicy returns a *ServerClientFilterConfig
// (session policy) if one can be found, otherwise nil is returned.
// Note that the connection is closed upon fatal error,
// when we fail to lookup the proc info, we return nil and an error.
func (s *ProxySession) getFilterPolicy() (*ServerClientFilterConfig, error) {

	fields := strings.Split(s.appConn.RemoteAddr().String(), ":")
	dstPortStr := fields[1]
	dstIP := net.ParseIP(s.cfg.ListenIP)

	if dstIP == nil {
		s.appConn.Close()
		return nil, fmt.Errorf("net.ParseIP fail for: %s\n", s.cfg.ListenIP)
	}

	srcP, _ := strconv.ParseUint(dstPortStr, 10, 16)
	dstP, _ := strconv.ParseUint(s.cfg.ListenTCPPort, 10, 16)
	procInfo := proc.LookupTCPSocketProcess(uint16(srcP), dstIP, uint16(dstP))
	log.Printf("proc exec path: %s\n", procInfo.ExePath)
	if procInfo == nil {
		s.appConn.Close()
		return nil, fmt.Errorf("Could not find process information for: %d %s %d\n", srcP, dstIP, dstP)
	}

	filter := getFilterForPathAndUID(procInfo.ExePath, procInfo.Uid)
	if filter == nil {
		filter = getFilterForPath(procInfo.ExePath)
	} else {
		log.Printf("No filters found for: %s (%d)\n", procInfo.ExePath, procInfo.Uid)
		filter = nil
	}
	return filter, nil
}

func (s *ProxySession) initTorControl() error {
	var err error

	// Connect to the real control port.
	if s.torConn, err = bulb.Dial("unix", s.cfg.TorControlSocketPath); err != nil {
		return fmt.Errorf("ERR/tor: Failed to connect to tor control port: %v", err)
	}

	// Issue a PROTOCOLINFO, so we can send a realistic response.
	if s.protoInfo, err = s.torConn.ProtocolInfo(); err != nil {
		s.torConn.Close()
		return fmt.Errorf("ERR/tor: Failed to issue PROTOCOLINFO: %v", err)
	}

	// Authenticate with the real tor control port.
	// XXX: Pull password out of `b.s.cfg`.
	if err = s.torConn.Authenticate(""); err != nil {
		s.torConn.Close()
		return fmt.Errorf("ERR/tor: Failed to authenticate: %v", err)
	}
	return nil
}

func (s *ProxySession) sessionWorker() {
	defer s.appConn.Close()

	var policy *ServerClientFilterConfig
	var err error

	clientAddr := s.appConn.RemoteAddr()
	log.Printf("INFO/tor: New ctrl connection from: %s", clientAddr)

	policy, err = s.getFilterPolicy()
	if err != nil {
		log.Printf("proc info query failure; connection from %s aborted: %s\n", clientAddr, err)
		return
	}
	if policy == nil {
		s.clientFilterPolicy = nil
		s.serverFilterPolicy = nil
	} else {
		s.clientFilterPolicy = &FilterConfig{
			Allowed:             policy.ClientAllowed,
			AllowedPrefixes:     policy.ClientAllowedPrefixes,
			Replacements:        policy.ClientReplacements,
			ReplacementPrefixes: policy.ClientReplacementPrefixes,
		}
		s.serverFilterPolicy = &FilterConfig{
			Allowed:             policy.ServerAllowed,
			AllowedPrefixes:     policy.ServerAllowedPrefixes,
			Replacements:        policy.ServerReplacements,
			ReplacementPrefixes: policy.ServerReplacementPrefixes,
		}
	}

	// Authenticate with the real control port
	err = s.initTorControl()
	defer s.torConn.Close()

	if err != nil {
		log.Printf("RoflcopTor: Failed to authenticate with the tor control port: %s\n", err)
		return
	}

	s.appConnReader = bufio.NewReader(s.appConn)
	s.Add(2)
	go s.proxyFilterTorToApp()
	go s.proxyFilterAppToTor()

	// Wait till all sessions are finished, log and return.
	s.Wait()
	s.torConn.Close()
	s.appConn.Close()

	if len(s.errChan) > 0 {
		err := <-s.errChan
		log.Printf("INFO/tor: Closed client connection from: %s: %v", clientAddr, err)
	} else {
		log.Printf("INFO/tor: Closed client connection from: %v", clientAddr)
	}
}

func (s *ProxySession) proxyFilterTorToApp() {
	defer s.Done()

	var appConnLock sync.Mutex
	writeAppConn := func(b []byte) (int, error) {
		appConnLock.Lock() // XXX is the lock really needed here?
		defer appConnLock.Unlock()
		return s.appConn.Write(b)
	}

	rd := bufio.NewReader(s.torConn)
	for {
		line, err := rd.ReadBytes('\n')
		if err != nil {
			s.errChan <- err
			break
		}
		lineStr := strings.TrimSpace(string(line))
		log.Printf("A<-T: [%s]\n", lineStr)
		if s.watch {
			if _, err = writeAppConn([]byte(lineStr + "\n")); err != nil {
				s.errChan <- err
				break
			}
		} else {
			filterCommand(lineStr, "250 OK", writeAppConn, s.errChan, s.serverFilterPolicy)
		}
	}
}

func (s *ProxySession) proxyFilterAppToTor() {
	defer s.Done()

	writeToTor := func(line []byte) (int, error) {
		n, err := s.torConn.Write([]byte(line))
		return n, err
	}

	for {
		line, err := s.appConnReader.ReadBytes('\n')
		if err != nil {
			s.errChan <- err
			break
		}
		lineStr := strings.TrimSpace(string(line))
		log.Printf("A->T: [%s]\n", lineStr)

		if s.watch {
			_, err = writeToTor([]byte(lineStr + "\n"))
			if err != nil {
				s.errChan <- err
			}
		} else {
			filterCommand(lineStr, "", writeToTor, s.errChan, s.clientFilterPolicy)
		}
	}
}
<|endoftext|>


<commit_before>package agentintegration_test

import (
	"os"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
	"github.com/pivotal-cf/cf-redis-broker/integration/helpers"
	"github.com/pivotal-cf/cf-redis-broker/redisconf"
)

var _ = Describe("Startup", func() {
	var agentSession *gexec.Session

	AfterEach(func() {
		stopAgent(agentSession)
	})

	Context("when there is no redis.conf", func() {
		BeforeEach(func() {
			os.Remove(redisConfPath)
			agentSession = startAgent()
		})

		loadRedisConfFileWhenItExists := func() redisconf.Conf {
			Eventually(func() bool {
				return helpers.FileExists(redisConfPath)
			}).Should(BeTrue())
			conf, err := redisconf.Load(redisConfPath)
			Expect(err).ToNot(HaveOccurred())
			return conf
		}

		It("copies a redis.conf from the default path and adds a password", func() {
			conf := loadRedisConfFileWhenItExists()

			Expect(conf.Get("daemonize")).To(Equal("no"))
			Expect(conf.HasKey("requirepass")).To(BeTrue())
		})

		It("creates a new password each time", func() {
			initialConf := loadRedisConfFileWhenItExists()

			helpers.KillProcess(agentSession)
			err := os.Remove(redisConfPath)
			Expect(err).ToNot(HaveOccurred())

			agentSession = startAgent()
			newConf := loadRedisConfFileWhenItExists()
			Expect(initialConf.Get("requirepass")).NotTo(Equal(newConf.Get("requirepass")))
		})
	})

	Context("when there is a redis.conf already", func() {
		BeforeEach(func() {
			err := redisconf.New(
				redisconf.Param{Key: "daemonize", Value: "yes"},
				redisconf.Param{Key: "requirepass", Value: "someotherpassword"},
				redisconf.Param{Key: "shouldbedeleted", Value: "yes"},
			).Save(redisConfPath)
			Expect(err).ToNot(HaveOccurred())
			agentSession = startAgent()
		})

		Describe("The copied redis.conf file", func() {
			It("has it's original password", func() {
				newRedisConf, err := redisconf.Load(redisConfPath)
				Expect(err).ToNot(HaveOccurred())
				Expect(newRedisConf.Get("requirepass")).To(Equal("someotherpassword"))
			})

			It("resets other parameters", func() {
				newRedisConf, err := redisconf.Load(redisConfPath)
				Expect(err).ToNot(HaveOccurred())
				Expect(newRedisConf.Get("daemonize")).To(Equal("no"))
			})

			It("does not have additional parameters", func() {
				newRedisConf, err := redisconf.Load(redisConfPath)
				Expect(err).ToNot(HaveOccurred())
				Expect(newRedisConf.HasKey("shouldbedeleted")).To(BeFalse())
			})
		})
	})
})
<commit_msg>Fix typo<commit_after>package agentintegration_test

import (
	"os"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
	"github.com/onsi/gomega/gexec"
	"github.com/pivotal-cf/cf-redis-broker/integration/helpers"
	"github.com/pivotal-cf/cf-redis-broker/redisconf"
)

var _ = Describe("Startup", func() {
	var agentSession *gexec.Session

	AfterEach(func() {
		stopAgent(agentSession)
	})

	Context("when there is no redis.conf", func() {
		BeforeEach(func() {
			os.Remove(redisConfPath)
			agentSession = startAgent()
		})

		loadRedisConfFileWhenItExists := func() redisconf.Conf {
			Eventually(func() bool {
				return helpers.FileExists(redisConfPath)
			}).Should(BeTrue())
			conf, err := redisconf.Load(redisConfPath)
			Expect(err).ToNot(HaveOccurred())
			return conf
		}

		It("copies a redis.conf from the default path and adds a password", func() {
			conf := loadRedisConfFileWhenItExists()

			Expect(conf.Get("daemonize")).To(Equal("no"))
			Expect(conf.HasKey("requirepass")).To(BeTrue())
		})

		It("creates a new password each time", func() {
			initialConf := loadRedisConfFileWhenItExists()

			helpers.KillProcess(agentSession)
			err := os.Remove(redisConfPath)
			Expect(err).ToNot(HaveOccurred())

			agentSession = startAgent()
			newConf := loadRedisConfFileWhenItExists()
			Expect(initialConf.Get("requirepass")).NotTo(Equal(newConf.Get("requirepass")))
		})
	})

	Context("when there is a redis.conf already", func() {
		BeforeEach(func() {
			err := redisconf.New(
				redisconf.Param{Key: "daemonize", Value: "yes"},
				redisconf.Param{Key: "requirepass", Value: "someotherpassword"},
				redisconf.Param{Key: "shouldbedeleted", Value: "yes"},
			).Save(redisConfPath)
			Expect(err).ToNot(HaveOccurred())
			agentSession = startAgent()
		})

		Describe("The copied redis.conf file", func() {
			It("has its original password", func() {
				newRedisConf, err := redisconf.Load(redisConfPath)
				Expect(err).ToNot(HaveOccurred())
				Expect(newRedisConf.Get("requirepass")).To(Equal("someotherpassword"))
			})

			It("resets other parameters", func() {
				newRedisConf, err := redisconf.Load(redisConfPath)
				Expect(err).ToNot(HaveOccurred())
				Expect(newRedisConf.Get("daemonize")).To(Equal("no"))
			})

			It("does not have additional parameters", func() {
				newRedisConf, err := redisconf.Load(redisConfPath)
				Expect(err).ToNot(HaveOccurred())
				Expect(newRedisConf.HasKey("shouldbedeleted")).To(BeFalse())
			})
		})
	})
})
<|endoftext|>


<commit_before>package moxxiConf

import (
	"github.com/dchest/uniuri"
	"net"
	"os"
	"strings"
	"text/template"
)

func inArr(a []string, t string) bool {
	for _, s := range a {
		if t == s {
			return true
		}
	}
	return false
}

func validHost(s string) string {
	s = strings.Trim(s, ".")
	parts := strings.Split(s, DomainSep)
	if len(parts) < 2 {
		return ""
	}
	for i := 0; i < len(parts)-1; {
		if len(parts[i]) < 1 {
			parts = append(parts[:i], parts[i+1:]...)
		} else {
			i++
		}
	}
	return strings.Join(parts, DomainSep)
}

func confCheck(host, ip string, destTLS bool, port int, blockedHeaders []string) (siteParams, error) {
	var conf siteParams
	if conf.IntHost = validHost(host); conf.IntHost == "" {
		return siteParams{}, &Err{Code: ErrBadHost, value: host}
	}

	tempIP := net.ParseIP(ip)
	if tempIP == nil {
		return siteParams{}, &Err{Code: ErrBadIP, value: ip}
	}

	conf.IntPort = 80
	if port > 0 && port < MaxAllowedPort {
		conf.IntPort = port
	}

	conf.IntIP = tempIP.String()
	conf.Encrypted = destTLS
	conf.StripHeaders = blockedHeaders

	return conf, nil
}

func confWrite(confPath, confExt, baseURL string, subdomainLen int, t template.Template,
	excludes []string) func(siteParams) (siteParams, error) {

	return func(config siteParams) (siteParams, error) {

		err := os.ErrExist
		var randPart, fileName string
		var f *os.File

		for randPart == "" || os.IsExist(err) {
			randPart = uniuri.NewLenChars(subdomainLen, SubdomainChars) + "." + baseURL
			// pick again
			if inArr(excludes, randPart) {
				continue
			}
			fileName = strings.TrimRight(confPath, PathSep) + PathSep
			fileName += randPart + DomainSep + strings.TrimLeft(confExt, DomainSep)
			f, err = os.Create(fileName)
		}

		config.ExtHost = randPart

		if err == os.ErrPermission {
			return siteParams{}, &Err{Code: ErrFilePerm, value: fileName, deepErr: err}
		} else if err != nil {
			return siteParams{}, &Err{Code: ErrFileUnexpect, value: fileName, deepErr: err}
		}

		tErr := t.Execute(f, config)

		if err = f.Close(); err != nil {
			return siteParams{}, &Err{Code: ErrCloseFile, value: fileName, deepErr: err}
		}

		if tErr != nil {
			if err = os.Remove(fileName); err != nil {
				return siteParams{}, &Err{Code: ErrRemoveFile, value: fileName, deepErr: err}
			}
		}

		return config, nil
	}
}
<commit_msg>added a minimum check to the random part bit<commit_after>package moxxiConf

import (
	"github.com/dchest/uniuri"
	"net"
	"os"
	"strings"
	"text/template"
)

func inArr(a []string, t string) bool {
	for _, s := range a {
		if t == s {
			return true
		}
	}
	return false
}

func validHost(s string) string {
	s = strings.Trim(s, ".")
	parts := strings.Split(s, DomainSep)
	if len(parts) < 2 {
		return ""
	}
	for i := 0; i < len(parts)-1; {
		if len(parts[i]) < 1 {
			parts = append(parts[:i], parts[i+1:]...)
		} else {
			i++
		}
	}
	return strings.Join(parts, DomainSep)
}

func confCheck(host, ip string, destTLS bool, port int, blockedHeaders []string) (siteParams, error) {
	var conf siteParams
	if conf.IntHost = validHost(host); conf.IntHost == "" {
		return siteParams{}, &Err{Code: ErrBadHost, value: host}
	}

	tempIP := net.ParseIP(ip)
	if tempIP == nil {
		return siteParams{}, &Err{Code: ErrBadIP, value: ip}
	}

	conf.IntPort = 80
	if port > 0 && port < MaxAllowedPort {
		conf.IntPort = port
	}

	conf.IntIP = tempIP.String()
	conf.Encrypted = destTLS
	conf.StripHeaders = blockedHeaders

	return conf, nil
}

func confWrite(confPath, confExt, baseURL string, subdomainLen int, t template.Template,
	excludes []string) func(siteParams) (siteParams, error) {

	if subDomainLen < 1 {
		subdomainLen = 1
	}

	return func(config siteParams) (siteParams, error) {

		err := os.ErrExist
		var randPart, fileName string
		var f *os.File

		for randPart == "" || os.IsExist(err) {
			randPart = uniuri.NewLenChars(subdomainLen, SubdomainChars) + "." + baseURL
			// pick again
			if inArr(excludes, randPart) {
				continue
			}
			fileName = strings.TrimRight(confPath, PathSep) + PathSep
			fileName += randPart + DomainSep + strings.TrimLeft(confExt, DomainSep)
			f, err = os.Create(fileName)
		}

		config.ExtHost = randPart

		if err == os.ErrPermission {
			return siteParams{}, &Err{Code: ErrFilePerm, value: fileName, deepErr: err}
		} else if err != nil {
			return siteParams{}, &Err{Code: ErrFileUnexpect, value: fileName, deepErr: err}
		}

		tErr := t.Execute(f, config)

		if err = f.Close(); err != nil {
			return siteParams{}, &Err{Code: ErrCloseFile, value: fileName, deepErr: err}
		}

		if tErr != nil {
			if err = os.Remove(fileName); err != nil {
				return siteParams{}, &Err{Code: ErrRemoveFile, value: fileName, deepErr: err}
			}
		}

		return config, nil
	}
}
<|endoftext|>


<commit_before>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package os provides a platform-independent interface to operating system
// functionality. The design is Unix-like, although the error handling is
// Go-like; failing calls return values of type error rather than error numbers.
// Often, more information is available within the error. For example,
// if a call that takes a file name fails, such as Open or Stat, the error
// will include the failing file name when printed and will be of type
// *PathError, which may be unpacked for more information.
//
// The os interface is intended to be uniform across all operating systems.
// Features not generally available appear in the system-specific package syscall.
//
// Here is a simple example, opening a file and reading some of it.
//
//	file, err := os.Open("file.go") // For read access.
//	if err != nil {
//		log.Fatal(err)
//	}
//
// If the open fails, the error string will be self-explanatory, like
//
//	open file.go: no such file or directory
//
// The file's data can then be read into a slice of bytes. Read and
// Write take their byte counts from the length of the argument slice.
//
//	data := make([]byte, 100)
//	count, err := file.Read(data)
//	if err != nil {
//		log.Fatal(err)
//	}
//	fmt.Printf("read %d bytes: %q\n", count, data[:count])
//
package os

import (
	"io"
	"syscall"
)

// Name returns the name of the file as presented to Open.
func (f *File) Name() string { return f.name }

// Stdin, Stdout, and Stderr are open Files pointing to the standard input,
// standard output, and standard error file descriptors.
//
// Note that the Go runtime writes to standard error for panics and crashes;
// closing Stderr may cause those messages to go elsewhere, perhaps
// to a file opened later.
var (
	Stdin  = NewFile(uintptr(syscall.Stdin), "/dev/stdin")
	Stdout = NewFile(uintptr(syscall.Stdout), "/dev/stdout")
	Stderr = NewFile(uintptr(syscall.Stderr), "/dev/stderr")
)

// Flags to OpenFile wrapping those of the underlying system. Not all
// flags may be implemented on a given system.
const (
	O_RDONLY int = syscall.O_RDONLY // open the file read-only.
	O_WRONLY int = syscall.O_WRONLY // open the file write-only.
	O_RDWR   int = syscall.O_RDWR   // open the file read-write.
	O_APPEND int = syscall.O_APPEND // append data to the file when writing.
	O_CREATE int = syscall.O_CREAT  // create a new file if none exists.
	O_EXCL   int = syscall.O_EXCL   // used with O_CREATE, file must not exist
	O_SYNC   int = syscall.O_SYNC   // open for synchronous I/O.
	O_TRUNC  int = syscall.O_TRUNC  // if possible, truncate file when opened.
)

// Seek whence values.
//
// Deprecated: Use io.SeekStart, io.SeekCurrent, and io.SeekEnd.
const (
	SEEK_SET int = 0 // seek relative to the origin of the file
	SEEK_CUR int = 1 // seek relative to the current offset
	SEEK_END int = 2 // seek relative to the end
)

// LinkError records an error during a link or symlink or rename
// system call and the paths that caused it.
type LinkError struct {
	Op  string
	Old string
	New string
	Err error
}

func (e *LinkError) Error() string {
	return e.Op + " " + e.Old + " " + e.New + ": " + e.Err.Error()
}

// Read reads up to len(b) bytes from the File.
// It returns the number of bytes read and any error encountered.
// At end of file, Read returns 0, io.EOF.
func (f *File) Read(b []byte) (n int, err error) {
	if err := f.checkValid("read"); err != nil {
		return 0, err
	}
	n, e := f.read(b)
	if n == 0 && len(b) > 0 && e == nil {
		return 0, io.EOF
	}
	if e != nil {
		err = &PathError{"read", f.name, e}
	}
	return n, err
}

// ReadAt reads len(b) bytes from the File starting at byte offset off.
// It returns the number of bytes read and the error, if any.
// ReadAt always returns a non-nil error when n < len(b).
// At end of file, that error is io.EOF.
func (f *File) ReadAt(b []byte, off int64) (n int, err error) {
	if err := f.checkValid("read"); err != nil {
		return 0, err
	}
	for len(b) > 0 {
		m, e := f.pread(b, off)
		if m == 0 && e == nil {
			return n, io.EOF
		}
		if e != nil {
			err = &PathError{"read", f.name, e}
			break
		}
		n += m
		b = b[m:]
		off += int64(m)
	}
	return
}

// Write writes len(b) bytes to the File.
// It returns the number of bytes written and an error, if any.
// Write returns a non-nil error when n != len(b).
func (f *File) Write(b []byte) (n int, err error) {
	if err := f.checkValid("write"); err != nil {
		return 0, err
	}
	n, e := f.write(b)
	if n < 0 {
		n = 0
	}
	if n != len(b) {
		err = io.ErrShortWrite
	}

	epipecheck(f, e)

	if e != nil {
		err = &PathError{"write", f.name, e}
	}
	return n, err
}

// WriteAt writes len(b) bytes to the File starting at byte offset off.
// It returns the number of bytes written and an error, if any.
// WriteAt returns a non-nil error when n != len(b).
func (f *File) WriteAt(b []byte, off int64) (n int, err error) {
	if err := f.checkValid("write"); err != nil {
		return 0, err
	}
	for len(b) > 0 {
		m, e := f.pwrite(b, off)
		if e != nil {
			err = &PathError{"write", f.name, e}
			break
		}
		n += m
		b = b[m:]
		off += int64(m)
	}
	return
}

// Seek sets the offset for the next Read or Write on file to offset, interpreted
// according to whence: 0 means relative to the origin of the file, 1 means
// relative to the current offset, and 2 means relative to the end.
// It returns the new offset and an error, if any.
// The behavior of Seek on a file opened with O_APPEND is not specified.
func (f *File) Seek(offset int64, whence int) (ret int64, err error) {
	if err := f.checkValid("seek"); err != nil {
		return 0, err
	}
	r, e := f.seek(offset, whence)
	if e == nil && f.dirinfo != nil && r != 0 {
		e = syscall.EISDIR
	}
	if e != nil {
		return 0, &PathError{"seek", f.name, e}
	}
	return r, nil
}

// WriteString is like Write, but writes the contents of string s rather than
// a slice of bytes.
func (f *File) WriteString(s string) (n int, err error) {
	return f.Write([]byte(s))
}

// Mkdir creates a new directory with the specified name and permission bits.
// If there is an error, it will be of type *PathError.
func Mkdir(name string, perm FileMode) error {
	e := syscall.Mkdir(fixLongPath(name), syscallMode(perm))

	if e != nil {
		return &PathError{"mkdir", name, e}
	}

	// mkdir(2) itself won't handle the sticky bit on *BSD and Solaris
	if !supportsCreateWithStickyBit && perm&ModeSticky != 0 {
		Chmod(name, perm)
	}

	return nil
}

// Chdir changes the current working directory to the named directory.
// If there is an error, it will be of type *PathError.
func Chdir(dir string) error {
	if e := syscall.Chdir(dir); e != nil {
		return &PathError{"chdir", dir, e}
	}
	return nil
}

// Chdir changes the current working directory to the file,
// which must be a directory.
// If there is an error, it will be of type *PathError.
func (f *File) Chdir() error {
	if err := f.checkValid("chdir"); err != nil {
		return err
	}
	if e := syscall.Fchdir(f.fd); e != nil {
		return &PathError{"chdir", f.name, e}
	}
	return nil
}

// Open opens the named file for reading. If successful, methods on
// the returned file can be used for reading; the associated file
// descriptor has mode O_RDONLY.
// If there is an error, it will be of type *PathError.
func Open(name string) (*File, error) {
	return OpenFile(name, O_RDONLY, 0)
}

// Create creates the named file with mode 0666 (before umask), truncating
// it if it already exists. If successful, methods on the returned
// File can be used for I/O; the associated file descriptor has mode
// O_RDWR.
// If there is an error, it will be of type *PathError.
func Create(name string) (*File, error) {
	return OpenFile(name, O_RDWR|O_CREATE|O_TRUNC, 0666)
}

// lstat is overridden in tests.
var lstat = Lstat

// Rename renames (moves) oldpath to newpath.
// If newpath already exists, Rename replaces it.
// OS-specific restrictions may apply when oldpath and newpath are in different directories.
// If there is an error, it will be of type *LinkError.
func Rename(oldpath, newpath string) error {
	return rename(oldpath, newpath)
}

// Many functions in package syscall return a count of -1 instead of 0.
// Using fixCount(call()) instead of call() corrects the count.
func fixCount(n int, err error) (int, error) {
	if n < 0 {
		n = 0
	}
	return n, err
}

// checkValid checks whether f is valid for use.
// If not, it returns an appropriate error, perhaps incorporating the operation name op.
func (f *File) checkValid(op string) error {
	if f == nil {
		return ErrInvalid
	}
	if f.fd == badFd {
		return &PathError{op, f.name, ErrClosed}
	}
	return nil
}
<commit_msg>os: document that Rename does not replace an existing directory<commit_after>// Copyright 2009 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package os provides a platform-independent interface to operating system
// functionality. The design is Unix-like, although the error handling is
// Go-like; failing calls return values of type error rather than error numbers.
// Often, more information is available within the error. For example,
// if a call that takes a file name fails, such as Open or Stat, the error
// will include the failing file name when printed and will be of type
// *PathError, which may be unpacked for more information.
//
// The os interface is intended to be uniform across all operating systems.
// Features not generally available appear in the system-specific package syscall.
//
// Here is a simple example, opening a file and reading some of it.
//
//	file, err := os.Open("file.go") // For read access.
//	if err != nil {
//		log.Fatal(err)
//	}
//
// If the open fails, the error string will be self-explanatory, like
//
//	open file.go: no such file or directory
//
// The file's data can then be read into a slice of bytes. Read and
// Write take their byte counts from the length of the argument slice.
//
//	data := make([]byte, 100)
//	count, err := file.Read(data)
//	if err != nil {
//		log.Fatal(err)
//	}
//	fmt.Printf("read %d bytes: %q\n", count, data[:count])
//
package os

import (
	"io"
	"syscall"
)

// Name returns the name of the file as presented to Open.
func (f *File) Name() string { return f.name }

// Stdin, Stdout, and Stderr are open Files pointing to the standard input,
// standard output, and standard error file descriptors.
//
// Note that the Go runtime writes to standard error for panics and crashes;
// closing Stderr may cause those messages to go elsewhere, perhaps
// to a file opened later.
var (
	Stdin  = NewFile(uintptr(syscall.Stdin), "/dev/stdin")
	Stdout = NewFile(uintptr(syscall.Stdout), "/dev/stdout")
	Stderr = NewFile(uintptr(syscall.Stderr), "/dev/stderr")
)

// Flags to OpenFile wrapping those of the underlying system. Not all
// flags may be implemented on a given system.
const (
	O_RDONLY int = syscall.O_RDONLY // open the file read-only.
	O_WRONLY int = syscall.O_WRONLY // open the file write-only.
	O_RDWR   int = syscall.O_RDWR   // open the file read-write.
	O_APPEND int = syscall.O_APPEND // append data to the file when writing.
	O_CREATE int = syscall.O_CREAT  // create a new file if none exists.
	O_EXCL   int = syscall.O_EXCL   // used with O_CREATE, file must not exist
	O_SYNC   int = syscall.O_SYNC   // open for synchronous I/O.
	O_TRUNC  int = syscall.O_TRUNC  // if possible, truncate file when opened.
)

// Seek whence values.
//
// Deprecated: Use io.SeekStart, io.SeekCurrent, and io.SeekEnd.
const (
	SEEK_SET int = 0 // seek relative to the origin of the file
	SEEK_CUR int = 1 // seek relative to the current offset
	SEEK_END int = 2 // seek relative to the end
)

// LinkError records an error during a link or symlink or rename
// system call and the paths that caused it.
type LinkError struct {
	Op  string
	Old string
	New string
	Err error
}

func (e *LinkError) Error() string {
	return e.Op + " " + e.Old + " " + e.New + ": " + e.Err.Error()
}

// Read reads up to len(b) bytes from the File.
// It returns the number of bytes read and any error encountered.
// At end of file, Read returns 0, io.EOF.
func (f *File) Read(b []byte) (n int, err error) {
	if err := f.checkValid("read"); err != nil {
		return 0, err
	}
	n, e := f.read(b)
	if n == 0 && len(b) > 0 && e == nil {
		return 0, io.EOF
	}
	if e != nil {
		err = &PathError{"read", f.name, e}
	}
	return n, err
}

// ReadAt reads len(b) bytes from the File starting at byte offset off.
// It returns the number of bytes read and the error, if any.
// ReadAt always returns a non-nil error when n < len(b).
// At end of file, that error is io.EOF.
func (f *File) ReadAt(b []byte, off int64) (n int, err error) {
	if err := f.checkValid("read"); err != nil {
		return 0, err
	}
	for len(b) > 0 {
		m, e := f.pread(b, off)
		if m == 0 && e == nil {
			return n, io.EOF
		}
		if e != nil {
			err = &PathError{"read", f.name, e}
			break
		}
		n += m
		b = b[m:]
		off += int64(m)
	}
	return
}

// Write writes len(b) bytes to the File.
// It returns the number of bytes written and an error, if any.
// Write returns a non-nil error when n != len(b).
func (f *File) Write(b []byte) (n int, err error) {
	if err := f.checkValid("write"); err != nil {
		return 0, err
	}
	n, e := f.write(b)
	if n < 0 {
		n = 0
	}
	if n != len(b) {
		err = io.ErrShortWrite
	}

	epipecheck(f, e)

	if e != nil {
		err = &PathError{"write", f.name, e}
	}
	return n, err
}

// WriteAt writes len(b) bytes to the File starting at byte offset off.
// It returns the number of bytes written and an error, if any.
// WriteAt returns a non-nil error when n != len(b).
func (f *File) WriteAt(b []byte, off int64) (n int, err error) {
	if err := f.checkValid("write"); err != nil {
		return 0, err
	}
	for len(b) > 0 {
		m, e := f.pwrite(b, off)
		if e != nil {
			err = &PathError{"write", f.name, e}
			break
		}
		n += m
		b = b[m:]
		off += int64(m)
	}
	return
}

// Seek sets the offset for the next Read or Write on file to offset, interpreted
// according to whence: 0 means relative to the origin of the file, 1 means
// relative to the current offset, and 2 means relative to the end.
// It returns the new offset and an error, if any.
// The behavior of Seek on a file opened with O_APPEND is not specified.
func (f *File) Seek(offset int64, whence int) (ret int64, err error) {
	if err := f.checkValid("seek"); err != nil {
		return 0, err
	}
	r, e := f.seek(offset, whence)
	if e == nil && f.dirinfo != nil && r != 0 {
		e = syscall.EISDIR
	}
	if e != nil {
		return 0, &PathError{"seek", f.name, e}
	}
	return r, nil
}

// WriteString is like Write, but writes the contents of string s rather than
// a slice of bytes.
func (f *File) WriteString(s string) (n int, err error) {
	return f.Write([]byte(s))
}

// Mkdir creates a new directory with the specified name and permission bits.
// If there is an error, it will be of type *PathError.
func Mkdir(name string, perm FileMode) error {
	e := syscall.Mkdir(fixLongPath(name), syscallMode(perm))

	if e != nil {
		return &PathError{"mkdir", name, e}
	}

	// mkdir(2) itself won't handle the sticky bit on *BSD and Solaris
	if !supportsCreateWithStickyBit && perm&ModeSticky != 0 {
		Chmod(name, perm)
	}

	return nil
}

// Chdir changes the current working directory to the named directory.
// If there is an error, it will be of type *PathError.
func Chdir(dir string) error {
	if e := syscall.Chdir(dir); e != nil {
		return &PathError{"chdir", dir, e}
	}
	return nil
}

// Chdir changes the current working directory to the file,
// which must be a directory.
// If there is an error, it will be of type *PathError.
func (f *File) Chdir() error {
	if err := f.checkValid("chdir"); err != nil {
		return err
	}
	if e := syscall.Fchdir(f.fd); e != nil {
		return &PathError{"chdir", f.name, e}
	}
	return nil
}

// Open opens the named file for reading. If successful, methods on
// the returned file can be used for reading; the associated file
// descriptor has mode O_RDONLY.
// If there is an error, it will be of type *PathError.
func Open(name string) (*File, error) {
	return OpenFile(name, O_RDONLY, 0)
}

// Create creates the named file with mode 0666 (before umask), truncating
// it if it already exists. If successful, methods on the returned
// File can be used for I/O; the associated file descriptor has mode
// O_RDWR.
// If there is an error, it will be of type *PathError.
func Create(name string) (*File, error) {
	return OpenFile(name, O_RDWR|O_CREATE|O_TRUNC, 0666)
}

// lstat is overridden in tests.
var lstat = Lstat

// Rename renames (moves) oldpath to newpath.
// If newpath already exists and is not a directory, Rename replaces it.
// OS-specific restrictions may apply when oldpath and newpath are in different directories.
// If there is an error, it will be of type *LinkError.
func Rename(oldpath, newpath string) error {
	return rename(oldpath, newpath)
}

// Many functions in package syscall return a count of -1 instead of 0.
// Using fixCount(call()) instead of call() corrects the count.
func fixCount(n int, err error) (int, error) {
	if n < 0 {
		n = 0
	}
	return n, err
}

// checkValid checks whether f is valid for use.
// If not, it returns an appropriate error, perhaps incorporating the operation name op.
func (f *File) checkValid(op string) error {
	if f == nil {
		return ErrInvalid
	}
	if f.fd == badFd {
		return &PathError{op, f.name, ErrClosed}
	}
	return nil
}
<|endoftext|>


<commit_before>package kubernetes

import (
	"strings"
	"testing"

	"k8s.io/client-go/kubernetes/fake"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/pkg/apis/extensions/v1beta1"
)

func TestDetectTargetContainer_with_name(t *testing.T) {
	deployment := &v1beta1.Deployment{
		ObjectMeta: v1.ObjectMeta{
			Name:      "deployment",
			Namespace: "default",
		},
		Spec: v1beta1.DeploymentSpec{
			Template: v1.PodTemplateSpec{
				Spec: v1.PodSpec{
					Containers: []v1.Container{
						v1.Container{
							Name: "rails",
						},
					},
				},
			},
		},
	}

	client := &Client{}

	testcases := []struct {
		name      string
		expectErr bool
		errMsg    string
	}{
		{
			name:      "rails",
			expectErr: false,
		},
		{
			name:      "foobar",
			expectErr: true,
			errMsg:    `container "foobar" does not exist`,
		},
	}

	for _, tc := range testcases {
		got, err := client.DetectTargetContainer(deployment, tc.name)

		if tc.expectErr {
			if err == nil {
				t.Error("got no error")
				continue
			}

			if !strings.Contains(err.Error(), tc.errMsg) {
				t.Errorf("error %q does not contain %q", err.Error(), tc.errMsg)
			}
		} else {
			if err != nil {
				t.Errorf("got error: %s", err)
				continue
			}

			if got.Name != tc.name {
				t.Errorf("expected deployment: %q, got: %q", tc.name, got.Name)
			}
		}
	}
}

func TestDetectTargetContainer_without_name(t *testing.T) {
	testcases := []struct {
		deployment *v1beta1.Deployment
		expectErr  bool
		errMsg     string
	}{
		{
			deployment: &v1beta1.Deployment{
				ObjectMeta: v1.ObjectMeta{
					Name:      "deployment",
					Namespace: "default",
				},
				Spec: v1beta1.DeploymentSpec{
					Template: v1.PodTemplateSpec{
						Spec: v1.PodSpec{
							Containers: []v1.Container{
								v1.Container{
									Name: "rails",
								},
							},
						},
					},
				},
			},
			expectErr: false,
		},
		{
			deployment: &v1beta1.Deployment{
				ObjectMeta: v1.ObjectMeta{
					Name:      "deployment",
					Namespace: "default",
				},
				Spec: v1beta1.DeploymentSpec{
					Template: v1.PodTemplateSpec{
						Spec: v1.PodSpec{
							Containers: []v1.Container{
								v1.Container{
									Name: "rails",
								},
								v1.Container{
									Name: "foobar",
								},
							},
						},
					},
				},
			},
			expectErr: true,
			errMsg:    `multiple containers ["rails" "foobar"] found in deployment "deployment"`,
		},
	}

	name := ""

	for _, tc := range testcases {
		client := &Client{}

		_, err := client.DetectTargetContainer(tc.deployment, name)

		if tc.expectErr {
			if err == nil {
				t.Error("got no error")
				continue
			}

			if !strings.Contains(err.Error(), tc.errMsg) {
				t.Errorf("error %q does not contain %q", err.Error(), tc.errMsg)
			}
		} else {
			if err != nil {
				t.Errorf("got error: %s", err)
			}
		}
	}
}

func TestDetectTargetDeployment_with_name(t *testing.T) {
	deployment := &v1beta1.Deployment{
		ObjectMeta: v1.ObjectMeta{
			Name:      "deployment",
			Namespace: "default",
		},
	}

	clientset := fake.NewSimpleClientset(deployment)
	client := &Client{
		clientset: clientset,
	}

	testcases := []struct {
		name      string
		expectErr bool
		errMsg    string
	}{
		{
			name:      "deployment",
			expectErr: false,
		},
		{
			name:      "foobar",
			expectErr: true,
			errMsg:    `failed to retrieve Deployment "foobar"`,
		},
	}

	namespace := "default"

	for _, tc := range testcases {
		got, err := client.DetectTargetDeployment(namespace, tc.name)

		if tc.expectErr {
			if err == nil {
				t.Error("got no error")
				continue
			}

			if !strings.Contains(err.Error(), tc.errMsg) {
				t.Errorf("error %q does not contain %q", err.Error(), tc.errMsg)
			}
		} else {
			if err != nil {
				t.Errorf("got error: %s", err)
				continue
			}

			if got.Name != tc.name {
				t.Errorf("expected deployment: %q, got: %q", tc.name, got.Name)
			}
		}
	}
}

func TestDetectTargetDeployment_without_name(t *testing.T) {
	testcases := []struct {
		deployments []v1beta1.Deployment
		expectErr   bool
		errMsg      string
	}{
		{
			deployments: []v1beta1.Deployment{
				v1beta1.Deployment{
					ObjectMeta: v1.ObjectMeta{
						Name:      "deployment",
						Namespace: "default",
					},
				},
			},
			expectErr: false,
		},
		{
			deployments: []v1beta1.Deployment{},
			expectErr:   true,
			errMsg:      `no Deployment found in namespace "default"`,
		},
		{
			deployments: []v1beta1.Deployment{
				v1beta1.Deployment{
					ObjectMeta: v1.ObjectMeta{
						Name:      "deployment",
						Namespace: "default",
					},
				},
				v1beta1.Deployment{
					ObjectMeta: v1.ObjectMeta{
						Name:      "foobar",
						Namespace: "default",
					},
				},
			},
			expectErr: true,
			errMsg:    `multiple Deployments ["deployment" "foobar"] found in namespace "default"`,
		},
	}

	name := ""
	namespace := "default"

	for _, tc := range testcases {
		clientset := fake.NewSimpleClientset(&v1beta1.DeploymentList{
			Items: tc.deployments,
		})
		client := &Client{
			clientset: clientset,
		}

		_, err := client.DetectTargetDeployment(namespace, name)

		if tc.expectErr {
			if err == nil {
				t.Error("got no error")
				continue
			}

			if !strings.Contains(err.Error(), tc.errMsg) {
				t.Errorf("error %q does not contain %q", err.Error(), tc.errMsg)
			}
		} else {
			if err != nil {
				t.Errorf("got error: %s", err)
			}
		}
	}
}

func TestGetDeployment(t *testing.T) {
	deployment := &v1beta1.Deployment{
		ObjectMeta: v1.ObjectMeta{
			Name:      "deployment",
			Namespace: "default",
		},
	}

	clientset := fake.NewSimpleClientset(deployment)
	client := &Client{
		clientset: clientset,
	}

	testcases := []struct {
		name      string
		expectErr bool
	}{
		{
			name:      "deployment",
			expectErr: false,
		},
		{
			name:      "foobar",
			expectErr: true,
		},
	}

	namespace := "default"

	for _, tc := range testcases {
		got, err := client.GetDeployment(namespace, tc.name)

		if tc.expectErr {
			if err == nil {
				t.Error("got no error")
			}
		} else {
			if err != nil {
				t.Errorf("got error: %s", err)
				continue
			}

			if got.Name != tc.name {
				t.Errorf("expected deployment: %q, got: %q", tc.name, got.Name)
			}
		}
	}
}

func TestSetImage(t *testing.T) {
	deployment := &v1beta1.Deployment{
		ObjectMeta: v1.ObjectMeta{
			Name:      "deployment",
			Namespace: "default",
		},
		Spec: v1beta1.DeploymentSpec{
			Template: v1.PodTemplateSpec{
				Spec: v1.PodSpec{
					Containers: []v1.Container{
						v1.Container{
							Name:  "rails",
							Image: "my-rails:v2",
						},
					},
				},
			},
		},
	}

	clientset := fake.NewSimpleClientset(deployment)
	client := &Client{
		clientset: clientset,
	}

	container := "rails"
	image := "my-rails:v3"

	got, err := client.SetImage(deployment, container, image)
	if err != nil {
		t.Errorf("got error: %s", err)
		return
	}

	gotImage := got.Spec.Template.Spec.Containers[0].Image
	if gotImage != image {
		t.Errorf("expected image: %q, got: %q", image, gotImage)
	}
}
<commit_msg>Add test for ListDeployment<commit_after>package kubernetes

import (
	"strings"
	"testing"

	"k8s.io/client-go/kubernetes/fake"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/pkg/apis/extensions/v1beta1"
)

func TestDetectTargetContainer_with_name(t *testing.T) {
	deployment := &v1beta1.Deployment{
		ObjectMeta: v1.ObjectMeta{
			Name:      "deployment",
			Namespace: "default",
		},
		Spec: v1beta1.DeploymentSpec{
			Template: v1.PodTemplateSpec{
				Spec: v1.PodSpec{
					Containers: []v1.Container{
						v1.Container{
							Name: "rails",
						},
					},
				},
			},
		},
	}

	client := &Client{}

	testcases := []struct {
		name      string
		expectErr bool
		errMsg    string
	}{
		{
			name:      "rails",
			expectErr: false,
		},
		{
			name:      "foobar",
			expectErr: true,
			errMsg:    `container "foobar" does not exist`,
		},
	}

	for _, tc := range testcases {
		got, err := client.DetectTargetContainer(deployment, tc.name)

		if tc.expectErr {
			if err == nil {
				t.Error("got no error")
				continue
			}

			if !strings.Contains(err.Error(), tc.errMsg) {
				t.Errorf("error %q does not contain %q", err.Error(), tc.errMsg)
			}
		} else {
			if err != nil {
				t.Errorf("got error: %s", err)
				continue
			}

			if got.Name != tc.name {
				t.Errorf("expected deployment: %q, got: %q", tc.name, got.Name)
			}
		}
	}
}

func TestDetectTargetContainer_without_name(t *testing.T) {
	testcases := []struct {
		deployment *v1beta1.Deployment
		expectErr  bool
		errMsg     string
	}{
		{
			deployment: &v1beta1.Deployment{
				ObjectMeta: v1.ObjectMeta{
					Name:      "deployment",
					Namespace: "default",
				},
				Spec: v1beta1.DeploymentSpec{
					Template: v1.PodTemplateSpec{
						Spec: v1.PodSpec{
							Containers: []v1.Container{
								v1.Container{
									Name: "rails",
								},
							},
						},
					},
				},
			},
			expectErr: false,
		},
		{
			deployment: &v1beta1.Deployment{
				ObjectMeta: v1.ObjectMeta{
					Name:      "deployment",
					Namespace: "default",
				},
				Spec: v1beta1.DeploymentSpec{
					Template: v1.PodTemplateSpec{
						Spec: v1.PodSpec{
							Containers: []v1.Container{
								v1.Container{
									Name: "rails",
								},
								v1.Container{
									Name: "foobar",
								},
							},
						},
					},
				},
			},
			expectErr: true,
			errMsg:    `multiple containers ["rails" "foobar"] found in deployment "deployment"`,
		},
	}

	name := ""

	for _, tc := range testcases {
		client := &Client{}

		_, err := client.DetectTargetContainer(tc.deployment, name)

		if tc.expectErr {
			if err == nil {
				t.Error("got no error")
				continue
			}

			if !strings.Contains(err.Error(), tc.errMsg) {
				t.Errorf("error %q does not contain %q", err.Error(), tc.errMsg)
			}
		} else {
			if err != nil {
				t.Errorf("got error: %s", err)
			}
		}
	}
}

func TestDetectTargetDeployment_with_name(t *testing.T) {
	deployment := &v1beta1.Deployment{
		ObjectMeta: v1.ObjectMeta{
			Name:      "deployment",
			Namespace: "default",
		},
	}

	clientset := fake.NewSimpleClientset(deployment)
	client := &Client{
		clientset: clientset,
	}

	testcases := []struct {
		name      string
		expectErr bool
		errMsg    string
	}{
		{
			name:      "deployment",
			expectErr: false,
		},
		{
			name:      "foobar",
			expectErr: true,
			errMsg:    `failed to retrieve Deployment "foobar"`,
		},
	}

	namespace := "default"

	for _, tc := range testcases {
		got, err := client.DetectTargetDeployment(namespace, tc.name)

		if tc.expectErr {
			if err == nil {
				t.Error("got no error")
				continue
			}

			if !strings.Contains(err.Error(), tc.errMsg) {
				t.Errorf("error %q does not contain %q", err.Error(), tc.errMsg)
			}
		} else {
			if err != nil {
				t.Errorf("got error: %s", err)
				continue
			}

			if got.Name != tc.name {
				t.Errorf("expected deployment: %q, got: %q", tc.name, got.Name)
			}
		}
	}
}

func TestDetectTargetDeployment_without_name(t *testing.T) {
	testcases := []struct {
		deployments []v1beta1.Deployment
		expectErr   bool
		errMsg      string
	}{
		{
			deployments: []v1beta1.Deployment{
				v1beta1.Deployment{
					ObjectMeta: v1.ObjectMeta{
						Name:      "deployment",
						Namespace: "default",
					},
				},
			},
			expectErr: false,
		},
		{
			deployments: []v1beta1.Deployment{},
			expectErr:   true,
			errMsg:      `no Deployment found in namespace "default"`,
		},
		{
			deployments: []v1beta1.Deployment{
				v1beta1.Deployment{
					ObjectMeta: v1.ObjectMeta{
						Name:      "deployment",
						Namespace: "default",
					},
				},
				v1beta1.Deployment{
					ObjectMeta: v1.ObjectMeta{
						Name:      "foobar",
						Namespace: "default",
					},
				},
			},
			expectErr: true,
			errMsg:    `multiple Deployments ["deployment" "foobar"] found in namespace "default"`,
		},
	}

	name := ""
	namespace := "default"

	for _, tc := range testcases {
		clientset := fake.NewSimpleClientset(&v1beta1.DeploymentList{
			Items: tc.deployments,
		})
		client := &Client{
			clientset: clientset,
		}

		_, err := client.DetectTargetDeployment(namespace, name)

		if tc.expectErr {
			if err == nil {
				t.Error("got no error")
				continue
			}

			if !strings.Contains(err.Error(), tc.errMsg) {
				t.Errorf("error %q does not contain %q", err.Error(), tc.errMsg)
			}
		} else {
			if err != nil {
				t.Errorf("got error: %s", err)
			}
		}
	}
}

func TestGetDeployment(t *testing.T) {
	deployment := &v1beta1.Deployment{
		ObjectMeta: v1.ObjectMeta{
			Name:      "deployment",
			Namespace: "default",
		},
	}

	clientset := fake.NewSimpleClientset(deployment)
	client := &Client{
		clientset: clientset,
	}

	testcases := []struct {
		name      string
		expectErr bool
	}{
		{
			name:      "deployment",
			expectErr: false,
		},
		{
			name:      "foobar",
			expectErr: true,
		},
	}

	namespace := "default"

	for _, tc := range testcases {
		got, err := client.GetDeployment(namespace, tc.name)

		if tc.expectErr {
			if err == nil {
				t.Error("got no error")
			}
		} else {
			if err != nil {
				t.Errorf("got error: %s", err)
				continue
			}

			if got.Name != tc.name {
				t.Errorf("expected deployment: %q, got: %q", tc.name, got.Name)
			}
		}
	}
}

func TestListDeployment(t *testing.T) {
	deployments := []v1beta1.Deployment{
		v1beta1.Deployment{
			ObjectMeta: v1.ObjectMeta{
				Name:      "deployment",
				Namespace: "default",
			},
		},
		v1beta1.Deployment{
			ObjectMeta: v1.ObjectMeta{
				Name:      "foobar",
				Namespace: "default",
			},
		},
	}

	clientset := fake.NewSimpleClientset(&v1beta1.DeploymentList{
		Items: deployments,
	})
	client := &Client{
		clientset: clientset,
	}

	namespace := "default"

	got, err := client.ListDeployment(namespace)
	if err != nil {
		t.Errorf("got error: %s", err)
	}

	expectedLength := 2
	if len(got) != expectedLength {
		t.Errorf("expected length: %d, got: %d", expectedLength, len(got))
	}
}

func TestSetImage(t *testing.T) {
	deployment := &v1beta1.Deployment{
		ObjectMeta: v1.ObjectMeta{
			Name:      "deployment",
			Namespace: "default",
		},
		Spec: v1beta1.DeploymentSpec{
			Template: v1.PodTemplateSpec{
				Spec: v1.PodSpec{
					Containers: []v1.Container{
						v1.Container{
							Name:  "rails",
							Image: "my-rails:v2",
						},
					},
				},
			},
		},
	}

	clientset := fake.NewSimpleClientset(deployment)
	client := &Client{
		clientset: clientset,
	}

	container := "rails"
	image := "my-rails:v3"

	got, err := client.SetImage(deployment, container, image)
	if err != nil {
		t.Errorf("got error: %s", err)
		return
	}

	gotImage := got.Spec.Template.Spec.Containers[0].Image
	if gotImage != image {
		t.Errorf("expected image: %q, got: %q", image, gotImage)
	}
}
<|endoftext|>


<commit_before>package modelhelper_test

import (
	"koding/db/models"
	"koding/db/mongodb/modelhelper"
	"koding/db/mongodb/modelhelper/modeltesthelper"
	"testing"
)

func TestSessionUpdateData(t *testing.T) {
	db := modeltesthelper.NewMongoDB(t)
	defer db.Close()
	testUsername := "testusername"
	testGroupName := "testgroupname"
	ses, err := modelhelper.CreateSessionForAccount(testUsername, testGroupName)
	if err != nil {
		t.Error(err)
	}

	nonExistingKey := "nonExistingKey"
	val, err := ses.Data.GetString(nonExistingKey)
	if err != models.ErrDataKeyNotExists {
		t.Error("expected ErrDataKeyNotExists, got", err)
	}

	if val != "" {
		t.Error("expected empty string, got", val)
	}

	key := "chargeID"
	value := "chargeVal"

	data := map[string]interface{}{
		key: value,
	}

	if err := modelhelper.UpdateSessionData(ses.ClientId, data); err != nil {
		t.Error(err)
	}

	ses, err = modelhelper.GetSessionById(ses.Id.Hex())
	if err != nil {
		t.Error(err)
	}

	val, err = ses.Data.GetString(key)
	if err != nil {
		t.Error("expected nil, got", err)
	}

	if val != value {
		t.Error("expected", value, "got", val)
	}
}
<commit_msg>modelhelper: add a test for GetMostRecentSession<commit_after>package modelhelper_test

import (
	"testing"
	"time"

	"koding/db/models"
	"koding/db/mongodb/modelhelper"
	"koding/db/mongodb/modelhelper/modeltesthelper"

	"gopkg.in/mgo.v2/bson"
)

func TestSessionUpdateData(t *testing.T) {
	db := modeltesthelper.NewMongoDB(t)
	defer db.Close()
	testUsername := "testusername"
	testGroupName := "testgroupname"
	ses, err := modelhelper.CreateSessionForAccount(testUsername, testGroupName)
	if err != nil {
		t.Error(err)
	}

	nonExistingKey := "nonExistingKey"
	val, err := ses.Data.GetString(nonExistingKey)
	if err != models.ErrDataKeyNotExists {
		t.Error("expected ErrDataKeyNotExists, got", err)
	}

	if val != "" {
		t.Error("expected empty string, got", val)
	}

	key := "chargeID"
	value := "chargeVal"

	data := map[string]interface{}{
		key: value,
	}

	if err := modelhelper.UpdateSessionData(ses.ClientId, data); err != nil {
		t.Error(err)
	}

	ses, err = modelhelper.GetSessionById(ses.Id.Hex())
	if err != nil {
		t.Error(err)
	}

	val, err = ses.Data.GetString(key)
	if err != nil {
		t.Error("expected nil, got", err)
	}

	if val != value {
		t.Error("expected", value, "got", val)
	}
}

func TestGetMostRecentSession(t *testing.T) {
	g, err := createGroup()
	if err != nil {
		t.Fatalf("createGroup()=%s", err)
	}

	newSession := func(long time.Duration) *models.Session {
		return &models.Session{
			Id:           bson.NewObjectId(),
			ClientId:     bson.NewObjectId().Hex(),
			Username:     "user",
			GroupName:    g.Slug,
			SessionBegan: time.Now().UTC().Add(long),
			LastAccess:   time.Now().UTC().Add(long),
		}
	}

	sessions := []*models.Session{
		newSession(0),
		newSession(10 * time.Minute),
		newSession(1 * time.Hour), // <- this is most recent
		newSession(15 * time.Second),
		newSession(58 * time.Minute),
		newSession(-1 * time.Hour),
	}

	for i, s := range sessions {
		if err := modelhelper.CreateSession(s); err != nil {
			t.Fatalf("%d: CreateSession()=%s", i, err)
		}
	}

	got, err := modelhelper.GetMostRecentSession("user")
	if err != nil {
		t.Fatalf("GetMostRecentSession()=%s", err)
	}

	if got.Id != sessions[2].Id {
		t.Fatalf("got %q, want %q", got.Id.Hex(), sessions[2].Id.Hex())
	}
}
<|endoftext|>


<commit_before>package koding

import (
	"koding/db/models"
	"koding/kites/kloud/contexthelper/session"
	"koding/kites/kloud/machinestate"
	"time"

	"github.com/koding/logging"
	"labix.org/v2/mgo/bson"
)

// Machine represents a single MongodDB document from the jMachines
// collection.
type Machine struct {
	Id          bson.ObjectId `bson:"_id" json:"-"`
	Label       string        `bson:"label"`
	Domain      string        `bson:"domain"`
	QueryString string        `bson:"queryString"`
	IpAddress   string        `bson:"ipAddress"`
	Assignee    struct {
		InProgress bool      `bson:"inProgress"`
		AssignedAt time.Time `bson:"assignedAt"`
	} `bson:"assignee"`
	Status struct {
		State      string    `bson:"state"`
		Reason     string    `bson:"reason"`
		ModifiedAt time.Time `bson:"modifiedAt"`
	} `bson:"status"`
	Provider   string    `bson:"provider"`
	Credential string    `bson:"credential"`
	CreatedAt  time.Time `bson:"createdAt"`
	Meta       struct {
		AlwaysOn     bool   `bson:"alwaysOn"`
		InstanceId   string `bson:"instanceId"`
		InstanceType string `bson:"instance_type"`
		InstanceName string `bson:"instanceName"`
		Region       string `bson:"region"`
		StorageSize  int    `bson:"storage_size"`
		SourceAmi    string `bson:"source_ami"`
		SnapshotId   string `bson:"-"`
	} `bson:"meta"`
	Users  []models.Permissions `bson:"users"`
	Groups []models.Permissions `bson:"groups"`

	// internal fields, not availabile in MongoDB schema
	Username             string           `bson:"-"`
	User                 *models.User     `bson:"-"`
	Payment              *PaymentResponse `bson:"-"`
	Session              *session.Session `bson:"-"`
	Log                  logging.Logger   `bson:"-"`
	networkUsageEndpoint string           `bson:"-"`

	// cleanFuncs are a list of functions that are called when after a method
	// is finished
	cleanFuncs []func()
}

// Cleanup calls all cleanup functions and set the list to nil. Once called any
// other call will not have any effect.
func (m *Machine) Cleanup() {
	if m.cleanFuncs == nil {
		return
	}

	for _, fn := range m.cleanFuncs {
		fn()
	}

	m.cleanFuncs = nil
}

func (m *Machine) State() machinestate.State {
	return machinestate.States[m.Status.State]
}
<commit_msg>kloud: be specific of how we are going to convert to map<commit_after>package koding

import (
	"koding/db/models"
	"koding/kites/kloud/contexthelper/session"
	"koding/kites/kloud/machinestate"
	"time"

	"github.com/koding/logging"
	"labix.org/v2/mgo/bson"
)

// Machine represents a single MongodDB document from the jMachines
// collection.
type Machine struct {
	Id          bson.ObjectId `bson:"_id" json:"-"`
	Label       string        `bson:"label"`
	Domain      string        `bson:"domain"`
	QueryString string        `bson:"queryString"`
	IpAddress   string        `bson:"ipAddress"`
	Assignee    struct {
		InProgress bool      `bson:"inProgress"`
		AssignedAt time.Time `bson:"assignedAt"`
	} `bson:"assignee"`
	Status struct {
		State      string    `bson:"state"`
		Reason     string    `bson:"reason"`
		ModifiedAt time.Time `bson:"modifiedAt"`
	} `bson:"status"`
	Provider   string    `bson:"provider"`
	Credential string    `bson:"credential"`
	CreatedAt  time.Time `bson:"createdAt"`
	Meta       struct {
		AlwaysOn     bool   `bson:"alwaysOn"`
		InstanceId   string `structs:"instanceId" bson:"instanceId"`
		InstanceType string `structs:"instance_type" bson:"instance_type"`
		InstanceName string `structs:"instanceName" bson:"instanceName"`
		Region       string `structs:"region" bson:"region"`
		StorageSize  int    `structs:"storage_size" bson:"storage_size"`
		SourceAmi    string `structs:"source_ami" bson:"source_ami"`
		SnapshotId   string `structs:"snapshotId" bson:"-"`
	} `bson:"meta"`
	Users  []models.Permissions `bson:"users"`
	Groups []models.Permissions `bson:"groups"`

	// internal fields, not availabile in MongoDB schema
	Username             string           `bson:"-"`
	User                 *models.User     `bson:"-"`
	Payment              *PaymentResponse `bson:"-"`
	Session              *session.Session `bson:"-"`
	Log                  logging.Logger   `bson:"-"`
	networkUsageEndpoint string           `bson:"-"`

	// cleanFuncs are a list of functions that are called when after a method
	// is finished
	cleanFuncs []func()
}

// Cleanup calls all cleanup functions and set the list to nil. Once called any
// other call will not have any effect.
func (m *Machine) Cleanup() {
	if m.cleanFuncs == nil {
		return
	}

	for _, fn := range m.cleanFuncs {
		fn()
	}

	m.cleanFuncs = nil
}

func (m *Machine) State() machinestate.State {
	return machinestate.States[m.Status.State]
}
<|endoftext|>


<commit_before>package collaboration

import (
	socialapimodels "socialapi/models"
	"socialapi/workers/collaboration/models"
	"strconv"
	"strings"
	"time"

	"koding/db/mongodb/modelhelper"
	"koding/kites/kloud/klient"

	"github.com/koding/bongo"
	"github.com/koding/kite"

	"labix.org/v2/mgo"
)

// DeleteDriveDoc deletes the file from google drive
func (c *Controller) DeleteDriveDoc(ping *models.Ping) error {
	// if file id is nil, there is nothing to do
	if ping.FileId == "" {
		return nil
	}

	return c.deleteFile(ping.FileId)
}

// EndPrivateMessage stops the collaboration session and deletes the all
// messages from db
func (c *Controller) EndPrivateMessage(ping *models.Ping) error {
	// if channel id is nil, there is nothing to do
	if ping.ChannelId == 0 {
		return nil
	}

	// fetch the channel
	channel := socialapimodels.NewChannel()
	if err := channel.ById(ping.ChannelId); err != nil {
		// if channel is not there, do not do anyting
		if err == bongo.RecordNotFound {
			return nil
		}

		return err
	}

	canOpen, err := channel.CanOpen(ping.AccountId)
	if err != nil {
		return err
	}

	if !canOpen {
		return nil // if the requester can not open the channel do not process
	}

	// delete the channel
	err = channel.Delete()
	if err != nil {
		return err
	}

	ws, err := modelhelper.GetWorkspaceByChannelId(
		strconv.FormatInt(ping.ChannelId, 10),
	)

	if err != nil && err != mgo.ErrNotFound {
		return err
	}

	// if the workspace is not there, nothing to do
	if err == mgo.ErrNotFound {
		return nil
	}

	return modelhelper.UnsetSocialChannelFromWorkspace(ws.ObjectId)
}

// UnshareVM removes the users from JMachine document
func (c *Controller) UnshareVM(ping *models.Ping) error {
	// if channel id is nil, there is nothing to do
	if ping.ChannelId == 0 {
		return nil
	}

	ws, err := modelhelper.GetWorkspaceByChannelId(
		strconv.FormatInt(ping.ChannelId, 10),
	)
	if err != nil && err != mgo.ErrNotFound {
		return err
	}

	// if the workspace is not there, nothing to do
	if err == mgo.ErrNotFound {
		return nil
	}

	return modelhelper.UnshareMachineByUid(ws.MachineUID)
}

// RemoveUsersFromMachine removes the collaboraters from the host machine
func (c *Controller) RemoveUsersFromMachine(ping *models.Ping) error {
	// if channel id is nil, there is nothing to do
	if ping.ChannelId == 0 {
		return nil
	}

	ws, err := modelhelper.GetWorkspaceByChannelId(strconv.FormatInt(ping.ChannelId, 10))
	if err != nil && err != mgo.ErrNotFound {
		return err
	}

	// if the workspace is not there, nothing to do
	if err == mgo.ErrNotFound {
		return nil
	}

	m, err := modelhelper.GetMachineByUid(ws.MachineUID)
	if err != nil && err != mgo.ErrNotFound {
		return err
	}

	// if the machine is not there, nothing to do
	if err == mgo.ErrNotFound {
		return nil
	}

	// Get the klient.
	klientRef, err := klient.ConnectTimeout(c.kite, m.QueryString, time.Second*10)
	if err != nil {
		if err == klient.ErrDialingFailed || err == kite.ErrNoKitesAvailable {
			c.log.Error(
				"[%s] Klient is not registered to Kontrol. Err: %s",
				m.QueryString,
				err,
			)

			return nil // if the machine is not open, we cant do anything
		}

		return err
	}
	defer klientRef.Close()

	type args struct {
		Username string

		// we are not gonna use this propery here, just for reference
		Permanent bool
	}

	var iterErr error
	for _, user := range m.Users {
		// do not unshare from owner user
		if user.Sudo && user.Owner {
			continue
		}

		// fetch user for its username
		u, err := modelhelper.GetUserById(user.Id.Hex())
		if err != nil {
			c.log.Error("couldnt get user", err.Error())

			// if we cant find the regarding user, do not do anything
			if err == mgo.ErrNotFound {
				continue
			}

			iterErr = err

			continue // do not stop iterating, unshare from others
		}

		param := args{
			Username: u.Name,
		}

		_, err = klientRef.Client.Tell("klient.unshare", param)
		if err != nil {
			c.log.Error("couldnt unshare %+v", err.Error())

			// those are so error prone, force klient side not to change the API
			// or make them exported to some other package?
			if strings.Contains(err.Error(), "user is permanent") {
				continue
			}

			if strings.Contains(err.Error(), "user is not in the shared list") {
				continue
			}

			if strings.Contains(err.Error(), "User not found") {
				continue
			}

			iterErr = err

			continue // do not stop iterating, unshare from others
		}
	}

	res, err := klientRef.Client.Tell("klient.shared", nil)
	if err == nil {
		c.log.Info("other users in the machine: %+v", res.MustString())
	}

	// iterErr will be nil if we dont encounter to any error in iter
	return iterErr

}
<commit_msg>Socialapi: fix the multiple workspace channel leave problems<commit_after>package collaboration

import (
	socialapimodels "socialapi/models"
	"socialapi/workers/collaboration/models"
	"strconv"
	"strings"
	"time"

	"koding/db/mongodb/modelhelper"
	"koding/kites/kloud/klient"

	"socialapi/request"

	"github.com/koding/bongo"
	"github.com/koding/kite"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

// DeleteDriveDoc deletes the file from google drive
func (c *Controller) DeleteDriveDoc(ping *models.Ping) error {
	// if file id is nil, there is nothing to do
	if ping.FileId == "" {
		return nil
	}

	return c.deleteFile(ping.FileId)
}

// EndPrivateMessage stops the collaboration session and deletes the all
// messages from db
func (c *Controller) EndPrivateMessage(ping *models.Ping) error {
	// if channel id is nil, there is nothing to do
	if ping.ChannelId == 0 {
		return nil
	}

	// fetch the channel
	channel := socialapimodels.NewChannel()
	if err := channel.ById(ping.ChannelId); err != nil {
		// if channel is not there, do not do anyting
		if err == bongo.RecordNotFound {
			return nil
		}

		return err
	}

	canOpen, err := channel.CanOpen(ping.AccountId)
	if err != nil {
		return err
	}

	if !canOpen {
		return nil // if the requester can not open the channel do not process
	}

	// delete the channel
	err = channel.Delete()
	if err != nil {
		return err
	}

	ws, err := modelhelper.GetWorkspaceByChannelId(
		strconv.FormatInt(ping.ChannelId, 10),
	)

	if err != nil && err != mgo.ErrNotFound {
		return err
	}

	// if the workspace is not there, nothing to do
	if err == mgo.ErrNotFound {
		return nil
	}

	return modelhelper.UnsetSocialChannelFromWorkspace(ws.ObjectId)
}

// UnshareVM removes the users from JMachine document
func (c *Controller) UnshareVM(ping *models.Ping) error {
	// if channel id is nil, there is nothing to do
	if ping.ChannelId == 0 {
		return nil
	}

	ws, err := modelhelper.GetWorkspaceByChannelId(
		strconv.FormatInt(ping.ChannelId, 10),
	)
	if err != nil && err != mgo.ErrNotFound {
		return err
	}

	// if the workspace is not there, nothing to do
	if err == mgo.ErrNotFound {
		return nil
	}

	toBeRemovedUsers, err := c.findToBeRemovedUsers(ping)
	if err != nil {
		return err
	}

	if len(toBeRemovedUsers) == 0 {
		return nil // no one to remove
	}

	return modelhelper.RemoveUsersFromMachineByIds(ws.MachineUID, toBeRemovedUsers)
}

func (c *Controller) findToBeRemovedUsers(ping *models.Ping) ([]bson.ObjectId, error) {
	ws, err := modelhelper.GetWorkspaceByChannelId(
		strconv.FormatInt(ping.ChannelId, 10),
	)
	if err != nil && err != mgo.ErrNotFound {
		return nil, err
	}

	// if the workspace is not there, nothing to do
	if err == mgo.ErrNotFound {
		return nil, nil
	}

	machine, err := modelhelper.GetMachineByUid(ws.MachineUID)
	if err != nil && err != mgo.ErrNotFound {
		return nil, err
	}

	if err == mgo.ErrNotFound {
		return nil, nil
	}

	ownerMachineUser := machine.Owner()
	if ownerMachineUser == nil {
		c.log.Critical("owner couldnt found %+v", ping)
		return nil, nil // if we cant find the owner, we cant process the users
	}

	ownerAccount, err := modelhelper.GetAccountByUserId(ownerMachineUser.Id)
	if err != nil && err != mgo.ErrNotFound {
		return nil, err
	}

	if err == mgo.ErrNotFound {
		return nil, nil
	}

	//	get workspaces of the owner
	ownersWorkspaces, err := modelhelper.GetWorkspaces(ownerMachineUser.Id)
	if err != nil && err != mgo.ErrNotFound {
		return nil, err
	}

	if err == mgo.ErrNotFound {
		return nil, nil
	}

	users := make(map[string]int)
	for _, ownerWS := range ownersWorkspaces {
		// if the workspace belongs to other vm, skip
		if ownerWS.MachineUID != ws.MachineUID {
			continue
		}

		channelId, err := strconv.ParseInt(ownerWS.ChannelId, 10, 64)
		if err != nil {
			continue
		}

		channel := socialapimodels.NewChannel()
		channel.Id = channelId
		participants, err := channel.FetchParticipants(&request.Query{})
		if err != nil {
			return nil, err
		}

		// fetch the channels of the workspaces
		for _, participant := range participants {
			// skip the owner
			if participant.OldId == ownerAccount.Id.Hex() {
				continue
			}

			if _, ok := users[participant.OldId]; !ok {
				users[participant.OldId] = 0
			}
			users[participant.OldId] = users[participant.OldId] + 1
		}
	}

	toBeRemovedUsers := make([]bson.ObjectId, 0)
	for accountId, count := range users {
		// if we count the user more than once, that means user is in another
		// workspace too
		if count > 1 {
			continue
		}

		u, err := modelhelper.GetUserByAccountId(accountId)
		if err != nil {
			return nil, err
		}

		toBeRemovedUsers = append(toBeRemovedUsers, u.ObjectId)
	}

	return toBeRemovedUsers, nil
}

// RemoveUsersFromMachine removes the collaboraters from the host machine
func (c *Controller) RemoveUsersFromMachine(ping *models.Ping) error {
	// if channel id is nil, there is nothing to do
	if ping.ChannelId == 0 {
		return nil
	}

	ws, err := modelhelper.GetWorkspaceByChannelId(strconv.FormatInt(ping.ChannelId, 10))
	if err != nil && err != mgo.ErrNotFound {
		return err
	}

	// if the workspace is not there, nothing to do
	if err == mgo.ErrNotFound {
		return nil
	}

	m, err := modelhelper.GetMachineByUid(ws.MachineUID)
	if err != nil && err != mgo.ErrNotFound {
		return err
	}

	// if the machine is not there, nothing to do
	if err == mgo.ErrNotFound {
		return nil
	}

	// Get the klient.
	klientRef, err := klient.ConnectTimeout(c.kite, m.QueryString, time.Second*10)
	if err != nil {
		if err == klient.ErrDialingFailed || err == kite.ErrNoKitesAvailable {
			c.log.Error(
				"[%s] Klient is not registered to Kontrol. Err: %s",
				m.QueryString,
				err,
			)

			return nil // if the machine is not open, we cant do anything
		}

		return err
	}
	defer klientRef.Close()

	type args struct {
		Username string

		// we are not gonna use this propery here, just for reference
		Permanent bool
	}

	var iterErr error
	for _, user := range m.Users {
		// do not unshare from owner user
		if user.Sudo && user.Owner {
			continue
		}

		// fetch user for its username
		u, err := modelhelper.GetUserById(user.Id.Hex())
		if err != nil {
			c.log.Error("couldnt get user", err.Error())

			// if we cant find the regarding user, do not do anything
			if err == mgo.ErrNotFound {
				continue
			}

			iterErr = err

			continue // do not stop iterating, unshare from others
		}

		param := args{
			Username: u.Name,
		}

		_, err = klientRef.Client.Tell("klient.unshare", param)
		if err != nil {
			c.log.Error("couldnt unshare %+v", err.Error())

			// those are so error prone, force klient side not to change the API
			// or make them exported to some other package?
			if strings.Contains(err.Error(), "user is permanent") {
				continue
			}

			if strings.Contains(err.Error(), "user is not in the shared list") {
				continue
			}

			if strings.Contains(err.Error(), "User not found") {
				continue
			}

			iterErr = err

			continue // do not stop iterating, unshare from others
		}
	}

	res, err := klientRef.Client.Tell("klient.shared", nil)
	if err == nil {
		c.log.Info("other users in the machine: %+v", res.MustString())
	}

	// iterErr will be nil if we dont encounter to any error in iter
	return iterErr

}
<|endoftext|>


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

package autoscale

import (
	"context"
	"fmt"

	"github.com/spf13/cobra"
	"k8s.io/klog/v2"

	autoscalingv1 "k8s.io/api/autoscaling/v1"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/cli-runtime/pkg/genericclioptions"
	"k8s.io/cli-runtime/pkg/printers"
	"k8s.io/cli-runtime/pkg/resource"
	autoscalingv1client "k8s.io/client-go/kubernetes/typed/autoscaling/v1"
	"k8s.io/client-go/scale"
	cmdutil "k8s.io/kubectl/pkg/cmd/util"
	"k8s.io/kubectl/pkg/scheme"
	"k8s.io/kubectl/pkg/util"
	"k8s.io/kubectl/pkg/util/i18n"
	"k8s.io/kubectl/pkg/util/templates"
)

var (
	autoscaleLong = templates.LongDesc(i18n.T(`
		Creates an autoscaler that automatically chooses and sets the number of pods that run in a kubernetes cluster.

		Looks up a Deployment, ReplicaSet, StatefulSet, or ReplicationController by name and creates an autoscaler that uses the given resource as a reference.
		An autoscaler can automatically increase or decrease number of pods deployed within the system as needed.`))

	autoscaleExample = templates.Examples(i18n.T(`
		# Auto scale a deployment "foo", with the number of pods between 2 and 10, no target CPU utilization specified so a default autoscaling policy will be used:
		kubectl autoscale deployment foo --min=2 --max=10

		# Auto scale a replication controller "foo", with the number of pods between 1 and 5, target CPU utilization at 80%:
		kubectl autoscale rc foo --max=5 --cpu-percent=80`))
)

// AutoscaleOptions declares the arguments accepted by the Autoscale command
type AutoscaleOptions struct {
	FilenameOptions *resource.FilenameOptions

	RecordFlags *genericclioptions.RecordFlags
	Recorder    genericclioptions.Recorder

	PrintFlags *genericclioptions.PrintFlags
	ToPrinter  func(string) (printers.ResourcePrinter, error)

	Name       string
	Min        int32
	Max        int32
	CPUPercent int32

	createAnnotation bool
	args             []string
	enforceNamespace bool
	namespace        string
	dryRunStrategy   cmdutil.DryRunStrategy
	dryRunVerifier   *resource.DryRunVerifier
	builder          *resource.Builder
	fieldManager     string

	HPAClient         autoscalingv1client.HorizontalPodAutoscalersGetter
	scaleKindResolver scale.ScaleKindResolver

	genericclioptions.IOStreams
}

// NewAutoscaleOptions creates the options for autoscale
func NewAutoscaleOptions(ioStreams genericclioptions.IOStreams) *AutoscaleOptions {
	return &AutoscaleOptions{
		PrintFlags:      genericclioptions.NewPrintFlags("autoscaled").WithTypeSetter(scheme.Scheme),
		FilenameOptions: &resource.FilenameOptions{},
		RecordFlags:     genericclioptions.NewRecordFlags(),
		Recorder:        genericclioptions.NoopRecorder{},

		IOStreams: ioStreams,
	}
}

// NewCmdAutoscale returns the autoscale Cobra command
func NewCmdAutoscale(f cmdutil.Factory, ioStreams genericclioptions.IOStreams) *cobra.Command {
	o := NewAutoscaleOptions(ioStreams)

	validArgs := []string{"deployment", "replicaset", "replicationcontroller", "statefulset"}

	cmd := &cobra.Command{
		Use:                   "autoscale (-f FILENAME | TYPE NAME | TYPE/NAME) [--min=MINPODS] --max=MAXPODS [--cpu-percent=CPU]",
		DisableFlagsInUseLine: true,
		Short:                 i18n.T("Auto-scale a Deployment, ReplicaSet, StatefulSet, or ReplicationController"),
		Long:                  autoscaleLong,
		Example:               autoscaleExample,
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(o.Complete(f, cmd, args))
			cmdutil.CheckErr(o.Validate())
			cmdutil.CheckErr(o.Run())
		},
		ValidArgs: validArgs,
	}

	// bind flag structs
	o.RecordFlags.AddFlags(cmd)
	o.PrintFlags.AddFlags(cmd)
	cmd.Flags().String("generator", "horizontalpodautoscaler/v1", i18n.T("The name of the API generator to use. Currently there is only 1 generator."))
	cmd.Flags().MarkDeprecated("generator", "has no effect and will be removed in the future.")
	cmd.Flags().Int32Var(&o.Min, "min", -1, "The lower limit for the number of pods that can be set by the autoscaler. If it's not specified or negative, the server will apply a default value.")
	cmd.Flags().Int32Var(&o.Max, "max", -1, "The upper limit for the number of pods that can be set by the autoscaler. Required.")
	cmd.MarkFlagRequired("max")
	cmd.Flags().Int32Var(&o.CPUPercent, "cpu-percent", -1, fmt.Sprintf("The target average CPU utilization (represented as a percent of requested CPU) over all the pods. If it's not specified or negative, a default autoscaling policy will be used."))
	cmd.Flags().StringVar(&o.Name, "name", "", i18n.T("The name for the newly created object. If not specified, the name of the input resource will be used."))
	cmdutil.AddDryRunFlag(cmd)
	cmdutil.AddFilenameOptionFlags(cmd, o.FilenameOptions, "identifying the resource to autoscale.")
	cmdutil.AddApplyAnnotationFlags(cmd)
	cmdutil.AddFieldManagerFlagVar(cmd, &o.fieldManager, "kubectl-autoscale")
	return cmd
}

// Complete verifies command line arguments and loads data from the command environment
func (o *AutoscaleOptions) Complete(f cmdutil.Factory, cmd *cobra.Command, args []string) error {
	var err error
	o.dryRunStrategy, err = cmdutil.GetDryRunStrategy(cmd)
	if err != nil {
		return err
	}
	dynamicClient, err := f.DynamicClient()
	if err != nil {
		return err
	}
	discoveryClient, err := f.ToDiscoveryClient()
	if err != nil {
		return err
	}
	o.dryRunVerifier = resource.NewDryRunVerifier(dynamicClient, f.OpenAPIGetter())
	o.createAnnotation = cmdutil.GetFlagBool(cmd, cmdutil.ApplyAnnotationsFlag)
	o.builder = f.NewBuilder()
	o.scaleKindResolver = scale.NewDiscoveryScaleKindResolver(discoveryClient)
	o.args = args
	o.RecordFlags.Complete(cmd)

	o.Recorder, err = o.RecordFlags.ToRecorder()
	if err != nil {
		return err
	}

	kubeClient, err := f.KubernetesClientSet()
	if err != nil {
		return err
	}
	o.HPAClient = kubeClient.AutoscalingV1()

	o.namespace, o.enforceNamespace, err = f.ToRawKubeConfigLoader().Namespace()
	if err != nil {
		return err
	}

	o.ToPrinter = func(operation string) (printers.ResourcePrinter, error) {
		o.PrintFlags.NamePrintFlags.Operation = operation
		cmdutil.PrintFlagsWithDryRunStrategy(o.PrintFlags, o.dryRunStrategy)

		return o.PrintFlags.ToPrinter()
	}

	return nil
}

// Validate checks that the provided attach options are specified.
func (o *AutoscaleOptions) Validate() error {
	if o.Max < 1 {
		return fmt.Errorf("--max=MAXPODS is required and must be at least 1, max: %d", o.Max)
	}
	if o.Max < o.Min {
		return fmt.Errorf("--max=MAXPODS must be larger or equal to --min=MINPODS, max: %d, min: %d", o.Max, o.Min)
	}

	return nil
}

// Run performs the execution
func (o *AutoscaleOptions) Run() error {
	r := o.builder.
		Unstructured().
		ContinueOnError().
		NamespaceParam(o.namespace).DefaultNamespace().
		FilenameParam(o.enforceNamespace, o.FilenameOptions).
		ResourceTypeOrNameArgs(false, o.args...).
		Flatten().
		Do()
	if err := r.Err(); err != nil {
		return err
	}

	count := 0
	err := r.Visit(func(info *resource.Info, err error) error {
		if err != nil {
			return err
		}

		mapping := info.ResourceMapping()
		gvr := mapping.GroupVersionKind.GroupVersion().WithResource(mapping.Resource.Resource)
		if _, err := o.scaleKindResolver.ScaleForResource(gvr); err != nil {
			return fmt.Errorf("cannot autoscale a %v: %v", mapping.GroupVersionKind.Kind, err)
		}

		hpa := o.createHorizontalPodAutoscaler(info.Name, mapping)

		if err := o.Recorder.Record(hpa); err != nil {
			klog.V(4).Infof("error recording current command: %v", err)
		}

		if o.dryRunStrategy == cmdutil.DryRunClient {
			count++

			printer, err := o.ToPrinter("created")
			if err != nil {
				return err
			}
			return printer.PrintObj(hpa, o.Out)
		}

		if err := util.CreateOrUpdateAnnotation(o.createAnnotation, hpa, scheme.DefaultJSONEncoder()); err != nil {
			return err
		}

		createOptions := metav1.CreateOptions{}
		if o.fieldManager != "" {
			createOptions.FieldManager = o.fieldManager
		}
		if o.dryRunStrategy == cmdutil.DryRunServer {
			if err := o.dryRunVerifier.HasSupport(hpa.GroupVersionKind()); err != nil {
				return err
			}
			createOptions.DryRun = []string{metav1.DryRunAll}
		}
		actualHPA, err := o.HPAClient.HorizontalPodAutoscalers(o.namespace).Create(context.TODO(), hpa, createOptions)
		if err != nil {
			return err
		}

		count++
		printer, err := o.ToPrinter("autoscaled")
		if err != nil {
			return err
		}
		return printer.PrintObj(actualHPA, o.Out)
	})
	if err != nil {
		return err
	}
	if count == 0 {
		return fmt.Errorf("no objects passed to autoscale")
	}
	return nil
}

func (o *AutoscaleOptions) createHorizontalPodAutoscaler(refName string, mapping *meta.RESTMapping) *autoscalingv1.HorizontalPodAutoscaler {
	name := o.Name
	if len(name) == 0 {
		name = refName
	}

	scaler := autoscalingv1.HorizontalPodAutoscaler{
		ObjectMeta: metav1.ObjectMeta{
			Name: name,
		},
		Spec: autoscalingv1.HorizontalPodAutoscalerSpec{
			ScaleTargetRef: autoscalingv1.CrossVersionObjectReference{
				APIVersion: mapping.GroupVersionKind.GroupVersion().String(),
				Kind:       mapping.GroupVersionKind.Kind,
				Name:       refName,
			},
			MaxReplicas: o.Max,
		},
	}

	if o.Min > 0 {
		v := int32(o.Min)
		scaler.Spec.MinReplicas = &v
	}
	if o.CPUPercent >= 0 {
		c := int32(o.CPUPercent)
		scaler.Spec.TargetCPUUtilizationPercentage = &c
	}

	return &scaler
}
<commit_msg>remove deprecated generator flag<commit_after>/*
Copyright 2015 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.
*/

package autoscale

import (
	"context"
	"fmt"

	"github.com/spf13/cobra"
	"k8s.io/klog/v2"

	autoscalingv1 "k8s.io/api/autoscaling/v1"
	"k8s.io/apimachinery/pkg/api/meta"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/cli-runtime/pkg/genericclioptions"
	"k8s.io/cli-runtime/pkg/printers"
	"k8s.io/cli-runtime/pkg/resource"
	autoscalingv1client "k8s.io/client-go/kubernetes/typed/autoscaling/v1"
	"k8s.io/client-go/scale"
	cmdutil "k8s.io/kubectl/pkg/cmd/util"
	"k8s.io/kubectl/pkg/scheme"
	"k8s.io/kubectl/pkg/util"
	"k8s.io/kubectl/pkg/util/i18n"
	"k8s.io/kubectl/pkg/util/templates"
)

var (
	autoscaleLong = templates.LongDesc(i18n.T(`
		Creates an autoscaler that automatically chooses and sets the number of pods that run in a kubernetes cluster.

		Looks up a Deployment, ReplicaSet, StatefulSet, or ReplicationController by name and creates an autoscaler that uses the given resource as a reference.
		An autoscaler can automatically increase or decrease number of pods deployed within the system as needed.`))

	autoscaleExample = templates.Examples(i18n.T(`
		# Auto scale a deployment "foo", with the number of pods between 2 and 10, no target CPU utilization specified so a default autoscaling policy will be used:
		kubectl autoscale deployment foo --min=2 --max=10

		# Auto scale a replication controller "foo", with the number of pods between 1 and 5, target CPU utilization at 80%:
		kubectl autoscale rc foo --max=5 --cpu-percent=80`))
)

// AutoscaleOptions declares the arguments accepted by the Autoscale command
type AutoscaleOptions struct {
	FilenameOptions *resource.FilenameOptions

	RecordFlags *genericclioptions.RecordFlags
	Recorder    genericclioptions.Recorder

	PrintFlags *genericclioptions.PrintFlags
	ToPrinter  func(string) (printers.ResourcePrinter, error)

	Name       string
	Min        int32
	Max        int32
	CPUPercent int32

	createAnnotation bool
	args             []string
	enforceNamespace bool
	namespace        string
	dryRunStrategy   cmdutil.DryRunStrategy
	dryRunVerifier   *resource.DryRunVerifier
	builder          *resource.Builder
	fieldManager     string

	HPAClient         autoscalingv1client.HorizontalPodAutoscalersGetter
	scaleKindResolver scale.ScaleKindResolver

	genericclioptions.IOStreams
}

// NewAutoscaleOptions creates the options for autoscale
func NewAutoscaleOptions(ioStreams genericclioptions.IOStreams) *AutoscaleOptions {
	return &AutoscaleOptions{
		PrintFlags:      genericclioptions.NewPrintFlags("autoscaled").WithTypeSetter(scheme.Scheme),
		FilenameOptions: &resource.FilenameOptions{},
		RecordFlags:     genericclioptions.NewRecordFlags(),
		Recorder:        genericclioptions.NoopRecorder{},

		IOStreams: ioStreams,
	}
}

// NewCmdAutoscale returns the autoscale Cobra command
func NewCmdAutoscale(f cmdutil.Factory, ioStreams genericclioptions.IOStreams) *cobra.Command {
	o := NewAutoscaleOptions(ioStreams)

	validArgs := []string{"deployment", "replicaset", "replicationcontroller", "statefulset"}

	cmd := &cobra.Command{
		Use:                   "autoscale (-f FILENAME | TYPE NAME | TYPE/NAME) [--min=MINPODS] --max=MAXPODS [--cpu-percent=CPU]",
		DisableFlagsInUseLine: true,
		Short:                 i18n.T("Auto-scale a Deployment, ReplicaSet, StatefulSet, or ReplicationController"),
		Long:                  autoscaleLong,
		Example:               autoscaleExample,
		Run: func(cmd *cobra.Command, args []string) {
			cmdutil.CheckErr(o.Complete(f, cmd, args))
			cmdutil.CheckErr(o.Validate())
			cmdutil.CheckErr(o.Run())
		},
		ValidArgs: validArgs,
	}

	// bind flag structs
	o.RecordFlags.AddFlags(cmd)
	o.PrintFlags.AddFlags(cmd)
	cmd.Flags().Int32Var(&o.Min, "min", -1, "The lower limit for the number of pods that can be set by the autoscaler. If it's not specified or negative, the server will apply a default value.")
	cmd.Flags().Int32Var(&o.Max, "max", -1, "The upper limit for the number of pods that can be set by the autoscaler. Required.")
	cmd.MarkFlagRequired("max")
	cmd.Flags().Int32Var(&o.CPUPercent, "cpu-percent", -1, fmt.Sprintf("The target average CPU utilization (represented as a percent of requested CPU) over all the pods. If it's not specified or negative, a default autoscaling policy will be used."))
	cmd.Flags().StringVar(&o.Name, "name", "", i18n.T("The name for the newly created object. If not specified, the name of the input resource will be used."))
	cmdutil.AddDryRunFlag(cmd)
	cmdutil.AddFilenameOptionFlags(cmd, o.FilenameOptions, "identifying the resource to autoscale.")
	cmdutil.AddApplyAnnotationFlags(cmd)
	cmdutil.AddFieldManagerFlagVar(cmd, &o.fieldManager, "kubectl-autoscale")
	return cmd
}

// Complete verifies command line arguments and loads data from the command environment
func (o *AutoscaleOptions) Complete(f cmdutil.Factory, cmd *cobra.Command, args []string) error {
	var err error
	o.dryRunStrategy, err = cmdutil.GetDryRunStrategy(cmd)
	if err != nil {
		return err
	}
	dynamicClient, err := f.DynamicClient()
	if err != nil {
		return err
	}
	discoveryClient, err := f.ToDiscoveryClient()
	if err != nil {
		return err
	}
	o.dryRunVerifier = resource.NewDryRunVerifier(dynamicClient, f.OpenAPIGetter())
	o.createAnnotation = cmdutil.GetFlagBool(cmd, cmdutil.ApplyAnnotationsFlag)
	o.builder = f.NewBuilder()
	o.scaleKindResolver = scale.NewDiscoveryScaleKindResolver(discoveryClient)
	o.args = args
	o.RecordFlags.Complete(cmd)

	o.Recorder, err = o.RecordFlags.ToRecorder()
	if err != nil {
		return err
	}

	kubeClient, err := f.KubernetesClientSet()
	if err != nil {
		return err
	}
	o.HPAClient = kubeClient.AutoscalingV1()

	o.namespace, o.enforceNamespace, err = f.ToRawKubeConfigLoader().Namespace()
	if err != nil {
		return err
	}

	o.ToPrinter = func(operation string) (printers.ResourcePrinter, error) {
		o.PrintFlags.NamePrintFlags.Operation = operation
		cmdutil.PrintFlagsWithDryRunStrategy(o.PrintFlags, o.dryRunStrategy)

		return o.PrintFlags.ToPrinter()
	}

	return nil
}

// Validate checks that the provided attach options are specified.
func (o *AutoscaleOptions) Validate() error {
	if o.Max < 1 {
		return fmt.Errorf("--max=MAXPODS is required and must be at least 1, max: %d", o.Max)
	}
	if o.Max < o.Min {
		return fmt.Errorf("--max=MAXPODS must be larger or equal to --min=MINPODS, max: %d, min: %d", o.Max, o.Min)
	}

	return nil
}

// Run performs the execution
func (o *AutoscaleOptions) Run() error {
	r := o.builder.
		Unstructured().
		ContinueOnError().
		NamespaceParam(o.namespace).DefaultNamespace().
		FilenameParam(o.enforceNamespace, o.FilenameOptions).
		ResourceTypeOrNameArgs(false, o.args...).
		Flatten().
		Do()
	if err := r.Err(); err != nil {
		return err
	}

	count := 0
	err := r.Visit(func(info *resource.Info, err error) error {
		if err != nil {
			return err
		}

		mapping := info.ResourceMapping()
		gvr := mapping.GroupVersionKind.GroupVersion().WithResource(mapping.Resource.Resource)
		if _, err := o.scaleKindResolver.ScaleForResource(gvr); err != nil {
			return fmt.Errorf("cannot autoscale a %v: %v", mapping.GroupVersionKind.Kind, err)
		}

		hpa := o.createHorizontalPodAutoscaler(info.Name, mapping)

		if err := o.Recorder.Record(hpa); err != nil {
			klog.V(4).Infof("error recording current command: %v", err)
		}

		if o.dryRunStrategy == cmdutil.DryRunClient {
			count++

			printer, err := o.ToPrinter("created")
			if err != nil {
				return err
			}
			return printer.PrintObj(hpa, o.Out)
		}

		if err := util.CreateOrUpdateAnnotation(o.createAnnotation, hpa, scheme.DefaultJSONEncoder()); err != nil {
			return err
		}

		createOptions := metav1.CreateOptions{}
		if o.fieldManager != "" {
			createOptions.FieldManager = o.fieldManager
		}
		if o.dryRunStrategy == cmdutil.DryRunServer {
			if err := o.dryRunVerifier.HasSupport(hpa.GroupVersionKind()); err != nil {
				return err
			}
			createOptions.DryRun = []string{metav1.DryRunAll}
		}
		actualHPA, err := o.HPAClient.HorizontalPodAutoscalers(o.namespace).Create(context.TODO(), hpa, createOptions)
		if err != nil {
			return err
		}

		count++
		printer, err := o.ToPrinter("autoscaled")
		if err != nil {
			return err
		}
		return printer.PrintObj(actualHPA, o.Out)
	})
	if err != nil {
		return err
	}
	if count == 0 {
		return fmt.Errorf("no objects passed to autoscale")
	}
	return nil
}

func (o *AutoscaleOptions) createHorizontalPodAutoscaler(refName string, mapping *meta.RESTMapping) *autoscalingv1.HorizontalPodAutoscaler {
	name := o.Name
	if len(name) == 0 {
		name = refName
	}

	scaler := autoscalingv1.HorizontalPodAutoscaler{
		ObjectMeta: metav1.ObjectMeta{
			Name: name,
		},
		Spec: autoscalingv1.HorizontalPodAutoscalerSpec{
			ScaleTargetRef: autoscalingv1.CrossVersionObjectReference{
				APIVersion: mapping.GroupVersionKind.GroupVersion().String(),
				Kind:       mapping.GroupVersionKind.Kind,
				Name:       refName,
			},
			MaxReplicas: o.Max,
		},
	}

	if o.Min > 0 {
		v := int32(o.Min)
		scaler.Spec.MinReplicas = &v
	}
	if o.CPUPercent >= 0 {
		c := int32(o.CPUPercent)
		scaler.Spec.TargetCPUUtilizationPercentage = &c
	}

	return &scaler
}
<|endoftext|>