<commit_before>package main

import (
	"crypto/sha256"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"strings"

	"github.com/lxc/lxd/shared"

	log "gopkg.in/inconshreveable/log15.v2"
)

const (
	APPARMOR_CMD_LOAD   = "r"
	APPARMOR_CMD_UNLOAD = "R"
	APPARMOR_CMD_PARSE  = "Q"
)

var aaPath = shared.VarPath("security", "apparmor")

const NESTING_AA_PROFILE = `
  pivot_root,
  mount /var/lib/lxd/shmounts/ -> /var/lib/lxd/shmounts/,
  mount none -> /var/lib/lxd/shmounts/,
  mount fstype=proc -> /usr/lib/x86_64-linux-gnu/lxc/**,
  mount fstype=sysfs -> /usr/lib/x86_64-linux-gnu/lxc/**,
  mount options=(rw,bind),
  mount options=(rw,rbind),
  deny /dev/.lxd/proc/** rw,
  deny /dev/.lxd/sys/** rw,
  mount options=(rw,make-rshared),

  # there doesn't seem to be a way to ask for:
  # mount options=(ro,nosuid,nodev,noexec,remount,bind),
  # as we always get mount to $cdir/proc/sys with those flags denied
  # So allow all mounts until that is straightened out:
  mount,
  mount options=bind /var/lib/lxd/shmounts/** -> /var/lib/lxd/**,
  change_profile -> lxc-container-default,
`

const DEFAULT_AA_PROFILE = `
#include <tunables/global>
profile "%s" flags=(attach_disconnected,mediate_deleted) {
    #include <abstractions/lxc/container-base>

    # user input raw.apparmor below here
    %s

    # nesting support goes here if needed
    %s
    change_profile -> "%s",
}`

func AAProfileFull(c *containerLXD) string {
	lxddir := shared.VarPath("")
	if len(c.name)+len(lxddir)+7 >= 253 {
		hash := sha256.New()
		io.WriteString(hash, lxddir)
		lxddir = fmt.Sprintf("%x", hash.Sum(nil))
	}

	return fmt.Sprintf("lxd-%s_<%s>", c.name, lxddir)
}

func AAProfileShort(c *containerLXD) string {
	return fmt.Sprintf("lxd-%s", c.name)
}

// getProfileContent generates the apparmor profile template from the given
// container. This includes the stock lxc includes as well as stuff from
// raw.apparmor.
func getAAProfileContent(c *containerLXD) string {
	rawApparmor, ok := c.config["raw.apparmor"]
	if !ok {
		rawApparmor = ""
	}

	nesting := ""
	if c.IsNesting() {
		nesting = NESTING_AA_PROFILE
	}

	return fmt.Sprintf(DEFAULT_AA_PROFILE, AAProfileFull(c), rawApparmor, nesting, AAProfileFull(c))
}

func runApparmor(command string, c *containerLXD) error {
	if aaConfined() {
		shared.Log.Debug("Already apparmor-confined (nested?), skipping aa profile actions")
		return nil
	}

	cmd := exec.Command("apparmor_parser", []string{
		fmt.Sprintf("-%sWL", command),
		path.Join(aaPath, "cache"),
		path.Join(aaPath, "profiles", AAProfileShort(c)),
	}...)

	output, err := cmd.CombinedOutput()
	if err != nil {
		shared.Log.Error("Running apparmor",
			log.Ctx{"action": command, "output": string(output), "err": err})
	}

	return err
}

/*
 * lxd could be confined by some other profile, but we'll only support
 * running under lxc-container-default-with-nesting or lxd-*
 */
func aaConfined() bool {
	curProfile := aaProfile()

	switch {
	case strings.HasPrefix(curProfile, "lxc-container-default-with-nesting"):
		return true
	case strings.HasPrefix(curProfile, "lxd-"):
		return true
	}

	return false
}

// Ensure that the container's policy is loaded into the kernel so the
// container can boot.
func AALoadProfile(c *containerLXD) error {
	if !aaEnabled {
		shared.Log.Debug("Apparmor not enabled, skipping profile load")
		return nil
	}

	/* In order to avoid forcing a profile parse (potentially slow) on
	 * every container start, let's use apparmor's binary policy cache,
	 * which checks mtime of the files to figure out if the policy needs to
	 * be regenerated.
	 *
	 * Since it uses mtimes, we shouldn't just always write out our local
	 * apparmor template; instead we should check to see whether the
	 * template is the same as ours. If it isn't we should write our
	 * version out so that the new changes are reflected and we definitely
	 * force a recompile.
	 */
	profile := path.Join(aaPath, "profiles", AAProfileShort(c))
	content, err := ioutil.ReadFile(profile)
	if err != nil && !os.IsNotExist(err) {
		return err
	}

	updated := getAAProfileContent(c)

	if string(content) != string(updated) {
		if err := os.MkdirAll(path.Join(aaPath, "profiles"), 0700); err != nil {
			return err
		}

		if err := ioutil.WriteFile(profile, []byte(updated), 0600); err != nil {
			return err
		}
	}

	return runApparmor(APPARMOR_CMD_LOAD, c)
}

// Ensure that the container's policy is unloaded to free kernel memory. This
// does not delete the policy from disk or cache.
func AAUnloadProfile(c *containerLXD) error {
	if !aaEnabled {
		shared.Log.Debug("Apparmor not enabled, skipping profile unload")
		return nil
	}

	return runApparmor(APPARMOR_CMD_UNLOAD, c)
}

// Parse the profile without loading it into the kernel.
func AAParseProfile(c *containerLXD) error {
	if !aaEnabled {
		shared.Log.Debug("Apparmor not enabled, skipping profile parse")
		return nil
	}

	return runApparmor(APPARMOR_CMD_PARSE, c)
}

// Delete the policy from cache/disk.
func AADeleteProfile(c *containerLXD) {
	if !aaEnabled {
		shared.Log.Debug("Apparmor not enabled, skipping profile delete")
		return
	}

	/* It's ok if these deletes fail: if the container was never started,
	 * we'll have never written a profile or cached it.
	 */
	os.Remove(path.Join(aaPath, "cache", AAProfileShort(c)))
	os.Remove(path.Join(aaPath, "profiles", AAProfileShort(c)))
}

// What's current apparmor profile
func aaProfile() string {
	contents, err := ioutil.ReadFile("/proc/self/attr/current")
	if err == nil {
		return strings.TrimSpace(string(contents))
	}
	return ""
}
<commit_msg>Add needed apparmor permissions for nesting<commit_after>package main

import (
	"crypto/sha256"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"strings"

	"github.com/lxc/lxd/shared"

	log "gopkg.in/inconshreveable/log15.v2"
)

const (
	APPARMOR_CMD_LOAD   = "r"
	APPARMOR_CMD_UNLOAD = "R"
	APPARMOR_CMD_PARSE  = "Q"
)

var aaPath = shared.VarPath("security", "apparmor")

const NESTING_AA_PROFILE = `
  pivot_root,
  mount /var/lib/lxd/shmounts/ -> /var/lib/lxd/shmounts/,
  mount none -> /var/lib/lxd/shmounts/,
  mount fstype=proc -> /usr/lib/x86_64-linux-gnu/lxc/**,
  mount fstype=sysfs -> /usr/lib/x86_64-linux-gnu/lxc/**,
  mount options=(rw,bind),
  mount options=(rw,rbind),
  deny /dev/.lxd/proc/** rw,
  deny /dev/.lxd/sys/** rw,
  mount options=(rw,make-rshared),

  # there doesn't seem to be a way to ask for:
  # mount options=(ro,nosuid,nodev,noexec,remount,bind),
  # as we always get mount to $cdir/proc/sys with those flags denied
  # So allow all mounts until that is straightened out:
  mount,
  mount options=bind /var/lib/lxd/shmounts/** -> /var/lib/lxd/**,
  change_profile -> lxc-container-default,
  # lxc-container-default-with-nesting also inherited these
  # from start-container, and seems to need them.
  ptrace,
  signal,
`

const DEFAULT_AA_PROFILE = `
#include <tunables/global>
profile "%s" flags=(attach_disconnected,mediate_deleted) {
    #include <abstractions/lxc/container-base>

    # user input raw.apparmor below here
    %s

    # nesting support goes here if needed
    %s
    change_profile -> "%s",
}`

func AAProfileFull(c *containerLXD) string {
	lxddir := shared.VarPath("")
	if len(c.name)+len(lxddir)+7 >= 253 {
		hash := sha256.New()
		io.WriteString(hash, lxddir)
		lxddir = fmt.Sprintf("%x", hash.Sum(nil))
	}

	return fmt.Sprintf("lxd-%s_<%s>", c.name, lxddir)
}

func AAProfileShort(c *containerLXD) string {
	return fmt.Sprintf("lxd-%s", c.name)
}

// getProfileContent generates the apparmor profile template from the given
// container. This includes the stock lxc includes as well as stuff from
// raw.apparmor.
func getAAProfileContent(c *containerLXD) string {
	rawApparmor, ok := c.config["raw.apparmor"]
	if !ok {
		rawApparmor = ""
	}

	nesting := ""
	if c.IsNesting() {
		nesting = NESTING_AA_PROFILE
	}

	return fmt.Sprintf(DEFAULT_AA_PROFILE, AAProfileFull(c), rawApparmor, nesting, AAProfileFull(c))
}

func runApparmor(command string, c *containerLXD) error {
	if aaConfined() {
		shared.Log.Debug("Already apparmor-confined (nested?), skipping aa profile actions")
		return nil
	}

	cmd := exec.Command("apparmor_parser", []string{
		fmt.Sprintf("-%sWL", command),
		path.Join(aaPath, "cache"),
		path.Join(aaPath, "profiles", AAProfileShort(c)),
	}...)

	output, err := cmd.CombinedOutput()
	if err != nil {
		shared.Log.Error("Running apparmor",
			log.Ctx{"action": command, "output": string(output), "err": err})
	}

	return err
}

/*
 * lxd could be confined by some other profile, but we'll only support
 * running under lxc-container-default-with-nesting or lxd-*
 */
func aaConfined() bool {
	curProfile := aaProfile()

	switch {
	case strings.HasPrefix(curProfile, "lxc-container-default-with-nesting"):
		return true
	case strings.HasPrefix(curProfile, "lxd-"):
		return true
	}

	return false
}

// Ensure that the container's policy is loaded into the kernel so the
// container can boot.
func AALoadProfile(c *containerLXD) error {
	if !aaEnabled {
		shared.Log.Debug("Apparmor not enabled, skipping profile load")
		return nil
	}

	/* In order to avoid forcing a profile parse (potentially slow) on
	 * every container start, let's use apparmor's binary policy cache,
	 * which checks mtime of the files to figure out if the policy needs to
	 * be regenerated.
	 *
	 * Since it uses mtimes, we shouldn't just always write out our local
	 * apparmor template; instead we should check to see whether the
	 * template is the same as ours. If it isn't we should write our
	 * version out so that the new changes are reflected and we definitely
	 * force a recompile.
	 */
	profile := path.Join(aaPath, "profiles", AAProfileShort(c))
	content, err := ioutil.ReadFile(profile)
	if err != nil && !os.IsNotExist(err) {
		return err
	}

	updated := getAAProfileContent(c)

	if string(content) != string(updated) {
		if err := os.MkdirAll(path.Join(aaPath, "profiles"), 0700); err != nil {
			return err
		}

		if err := ioutil.WriteFile(profile, []byte(updated), 0600); err != nil {
			return err
		}
	}

	return runApparmor(APPARMOR_CMD_LOAD, c)
}

// Ensure that the container's policy is unloaded to free kernel memory. This
// does not delete the policy from disk or cache.
func AAUnloadProfile(c *containerLXD) error {
	if !aaEnabled {
		shared.Log.Debug("Apparmor not enabled, skipping profile unload")
		return nil
	}

	return runApparmor(APPARMOR_CMD_UNLOAD, c)
}

// Parse the profile without loading it into the kernel.
func AAParseProfile(c *containerLXD) error {
	if !aaEnabled {
		shared.Log.Debug("Apparmor not enabled, skipping profile parse")
		return nil
	}

	return runApparmor(APPARMOR_CMD_PARSE, c)
}

// Delete the policy from cache/disk.
func AADeleteProfile(c *containerLXD) {
	if !aaEnabled {
		shared.Log.Debug("Apparmor not enabled, skipping profile delete")
		return
	}

	/* It's ok if these deletes fail: if the container was never started,
	 * we'll have never written a profile or cached it.
	 */
	os.Remove(path.Join(aaPath, "cache", AAProfileShort(c)))
	os.Remove(path.Join(aaPath, "profiles", AAProfileShort(c)))
}

// What's current apparmor profile
func aaProfile() string {
	contents, err := ioutil.ReadFile("/proc/self/attr/current")
	if err == nil {
		return strings.TrimSpace(string(contents))
	}
	return ""
}
<|endoftext|>


<commit_before>package lzma

import (
	"bufio"
	"errors"
	"io"
)

// states defines the overall state count
const states = 12

// bufferLen is the value used for the bufferLen used by the decoder.
var bufferLen = 64 * (1 << 10)

// Decoder is able to read a LZMA byte stream and to read the plain text.
type Decoder struct {
	properties       Properties
	unpackedLen      uint64
	decodedLen       uint64
	dict             *decoderDict
	state            uint32
	posBitMask       uint32
	rd               *rangeDecoder
	isMatch          [states << maxPosBits]prob
	isRep            [states]prob
	isRepG0          [states]prob
	isRepG1          [states]prob
	isRepG2          [states]prob
	isRepG0Long      [states << maxPosBits]prob
	rep              [4]uint32
	litDecoder       *literalCodec
	lengthDecoder    *lengthCodec
	repLengthDecoder *lengthCodec
	distDecoder      *distCodec
}

// NewDecoder creates an LZMA decoder. It reads the classic, original LZMA
// format. Note that LZMA2 uses a different header format.
func NewDecoder(r io.Reader) (d *Decoder, err error) {
	f := bufio.NewReader(r)
	properties, err := readProperties(f)
	if err != nil {
		return nil, err
	}
	historyLen := int(properties.DictLen)
	if historyLen < 0 {
		return nil, errors.New(
			"LZMA property DictLen exceeds maximum int value")
	}
	d = &Decoder{
		properties: *properties,
	}
	if d.unpackedLen, err = readUint64LE(f); err != nil {
		return nil, err
	}
	if d.dict, err = newDecoderDict(bufferLen, historyLen); err != nil {
		return nil, err
	}
	d.posBitMask = (uint32(1) << uint(d.properties.PB)) - 1
	if d.rd, err = newRangeDecoder(f); err != nil {
		return nil, err
	}
	initProbSlice(d.isMatch[:])
	initProbSlice(d.isRep[:])
	initProbSlice(d.isRepG0[:])
	initProbSlice(d.isRepG1[:])
	initProbSlice(d.isRepG2[:])
	initProbSlice(d.isRepG0Long[:])
	d.litDecoder = newLiteralCodec(d.properties.LC, d.properties.LP)
	d.lengthDecoder = newLengthCodec()
	d.repLengthDecoder = newLengthCodec()
	d.distDecoder = newDistCodec()
	return d, nil
}

// Properties returns a set of properties.
func (d *Decoder) Properties() Properties {
	return d.properties
}

// getUint64LE converts the uint64 value stored as little endian to an uint64
// value.
func getUint64LE(b []byte) uint64 {
	x := uint64(b[7]) << 56
	x |= uint64(b[6]) << 48
	x |= uint64(b[5]) << 40
	x |= uint64(b[4]) << 32
	x |= uint64(b[3]) << 24
	x |= uint64(b[2]) << 16
	x |= uint64(b[1]) << 8
	x |= uint64(b[0])
	return x
}

// readUint64LE reads a uint64 little-endian integer from reader.
func readUint64LE(r io.Reader) (x uint64, err error) {
	b := make([]byte, 8)
	if _, err = io.ReadFull(r, b); err != nil {
		return 0, err
	}
	x = getUint64LE(b)
	return x, nil
}

// initProbSlice initializes a slice of probabilities.
func initProbSlice(p []prob) {
	for i := range p {
		p[i] = probInit
	}
}

// Reads reads data from the decoder stream.
//
// The method might block and is not reentrant.
//
// The end of the LZMA stream is indicated by EOF. There might be other errors
// returned. The decoder will not be able to recover from an error returned.
func (d *Decoder) Read(p []byte) (n int, err error) {
	for {
		var k int
		k, err = d.dict.Read(p[n:])
		n += k
		switch {
		case err == io.EOF:
			if n <= 0 {
				return 0, io.EOF
			}
			return n, nil
		case err != nil:
			return n, err
		case n == len(p):
			return n, nil
		}
		if err = d.fill(); err != nil {
			return n, err
		}
	}
}

// fill puts at lest the requested number of bytes into the decoder dictionary.
func (d *Decoder) fill() error {
	for !d.dict.eof && d.dict.readable() < d.dict.b {
		op, err := d.decodeOp()
		if err != nil {
			switch {
			case err == eofDecoded:
				d.dict.eof = true
				return nil
			case err == io.EOF:
				return errors.New(
					"unexpected end of compressed stream")
			default:
				return err
			}
		}
		n := d.decodedLen + uint64(op.Len())
		if n < d.decodedLen {
			panic("negative op length or overflow of decodedLen")
		}
		d.decodedLen = n
		if d.decodedLen >= d.unpackedLen {
			if d.decodedLen == d.unpackedLen {
				d.dict.eof = true
				return nil
			}
			return errors.New(
				"decoded stream longer than provided length")
		}
		if err = op.applyDecoderDict(d.dict); err != nil {
			return err
		}
	}
	return nil
}

// updateStateLiteral updates the state for a literal.
func (d *Decoder) updateStateLiteral() {
	switch {
	case d.state < 4:
		d.state = 0
		return
	case d.state < 10:
		d.state -= 3
		return
	}
	d.state -= 6
	return
}

// updateStateMatch updates the state for a match.
func (d *Decoder) updateStateMatch() {
	if d.state < 7 {
		d.state = 7
		return
	}
	d.state = 10
	return
}

// updateStateRep updates the state for a repetition.
func (d *Decoder) updateStateRep() {
	if d.state < 7 {
		d.state = 8
	}
	d.state = 11
}

// updateStateShortRep updates the state for a short repetition.
func (d *Decoder) updateStateShortRep() {
	if d.state < 7 {
		d.state = 9
	}
	d.state = 11
}

// decodeLiteral decodes a literal.
func (d *Decoder) decodeLiteral() (op operation, err error) {
	prevByte := d.dict.getByte(1)
	lp, lc := uint(d.properties.LP), uint(d.properties.LC)
	litState := ((uint32(d.dict.total) & ((1 << lp) - 1)) << lc) |
		(uint32(prevByte) >> (8 - lc))
	match := d.dict.getByte(int(d.rep[0]) + 1)
	s, err := d.litDecoder.Decode(d.rd, d.state, match, litState)
	if err != nil {
		return nil, err
	}
	return lit{s}, nil
}

// errWrongTermination indicates that a termination symbol has been received,
// but the range decoder could still produces more data
var errWrongTermination = errors.New(
	"range decoder doesn't support termination")

// decodeOp indicates an EOF of the decoded file
var eofDecoded = errors.New("EOF of decoded stream")

// decodeOp decodes an operation. The function returns eofDecoded if there is
// an explicit termination marker.
func (d *Decoder) decodeOp() (op operation, err error) {
	posState := uint32(d.dict.total) & d.posBitMask
	state2 := (d.state << maxPosBits) | posState

	b, err := d.isMatch[state2].Decode(d.rd)
	if err != nil {
		return nil, err
	}
	if b == 0 {
		// literal
		op, err := d.decodeLiteral()
		if err != nil {
			return nil, err
		}
		d.updateStateLiteral()
		return op, nil
	}
	b, err = d.isRep[d.state].Decode(d.rd)
	if err != nil {
		return nil, err
	}
	if b == 0 {
		// simple match
		d.rep[3], d.rep[2], d.rep[1] = d.rep[2], d.rep[1], d.rep[0]
		d.updateStateMatch()
		// The length decoder returns the length offset.
		l, err := d.lengthDecoder.Decode(d.rd, posState)
		if err != nil {
			return nil, err
		}
		// The dist decoder returns the distance offset. The actual
		// distance is 1 higher.
		d.rep[0], err = d.distDecoder.Decode(l, d.rd)
		if err != nil {
			return nil, err
		}
		if d.rep[0] == 0xffffffff {
			if !d.rd.possiblyAtEnd() {
				return nil, errWrongTermination
			}
			return nil, eofDecoded
		}
		op = rep{length: int(l) + minLength,
			distance: int(d.rep[0]) + minDistance}
		return op, nil
	}
	b, err = d.isRepG0[d.state].Decode(d.rd)
	if err != nil {
		return nil, err
	}
	dist := d.rep[0]
	if b == 0 {
		// rep match 0
		b, err = d.isRepG0Long[state2].Decode(d.rd)
		if err != nil {
			return nil, err
		}
		if b == 0 {
			d.updateStateShortRep()
			op = rep{length: 1,
				distance: int(d.rep[0]) + minDistance}
			return op, nil
		}
	} else {
		b, err = d.isRepG1[d.state].Decode(d.rd)
		if err != nil {
			return nil, err
		}
		if b == 0 {
			dist = d.rep[1]
		} else {
			b, err = d.isRepG2[d.state].Decode(d.rd)
			if err != nil {
				return nil, err
			}
			if b == 0 {
				dist = d.rep[2]
			} else {
				dist = d.rep[3]
				d.rep[3] = d.rep[2]
			}
			d.rep[2] = d.rep[1]
		}
		d.rep[1] = d.rep[0]
		d.rep[0] = dist
	}
	l, err := d.repLengthDecoder.Decode(d.rd, posState)
	if err != nil {
		return nil, err
	}
	d.updateStateRep()
	op = rep{length: int(l) + minLength, distance: int(dist) + minDistance}
	return op, nil
}
<commit_msg>lzma: handle an undefined unpack length<commit_after>package lzma

import (
	"bufio"
	"errors"
	"io"
)

// states defines the overall state count
const states = 12

// bufferLen is the value used for the bufferLen used by the decoder.
var bufferLen = 64 * (1 << 10)

// noUnpackLen requires an explicit end of stream marker
const noUnpackLen uint64 = 1<<64 - 1

// Decoder is able to read a LZMA byte stream and to read the plain text.
//
// Note that an unpackLen of 0xffffffffffffffff requires an explicit end of
// stream marker.
type Decoder struct {
	properties Properties
	// length to unpack
	unpackLen        uint64
	decodedLen       uint64
	dict             *decoderDict
	state            uint32
	posBitMask       uint32
	rd               *rangeDecoder
	isMatch          [states << maxPosBits]prob
	isRep            [states]prob
	isRepG0          [states]prob
	isRepG1          [states]prob
	isRepG2          [states]prob
	isRepG0Long      [states << maxPosBits]prob
	rep              [4]uint32
	litDecoder       *literalCodec
	lengthDecoder    *lengthCodec
	repLengthDecoder *lengthCodec
	distDecoder      *distCodec
}

// NewDecoder creates an LZMA decoder. It reads the classic, original LZMA
// format. Note that LZMA2 uses a different header format.
func NewDecoder(r io.Reader) (d *Decoder, err error) {
	f := bufio.NewReader(r)
	properties, err := readProperties(f)
	if err != nil {
		return nil, err
	}
	historyLen := int(properties.DictLen)
	if historyLen < 0 {
		return nil, errors.New(
			"LZMA property DictLen exceeds maximum int value")
	}
	d = &Decoder{
		properties: *properties,
	}
	if d.unpackLen, err = readUint64LE(f); err != nil {
		return nil, err
	}
	if d.dict, err = newDecoderDict(bufferLen, historyLen); err != nil {
		return nil, err
	}
	d.posBitMask = (uint32(1) << uint(d.properties.PB)) - 1
	if d.rd, err = newRangeDecoder(f); err != nil {
		return nil, err
	}
	initProbSlice(d.isMatch[:])
	initProbSlice(d.isRep[:])
	initProbSlice(d.isRepG0[:])
	initProbSlice(d.isRepG1[:])
	initProbSlice(d.isRepG2[:])
	initProbSlice(d.isRepG0Long[:])
	d.litDecoder = newLiteralCodec(d.properties.LC, d.properties.LP)
	d.lengthDecoder = newLengthCodec()
	d.repLengthDecoder = newLengthCodec()
	d.distDecoder = newDistCodec()
	return d, nil
}

// Properties returns a set of properties.
func (d *Decoder) Properties() Properties {
	return d.properties
}

// getUint64LE converts the uint64 value stored as little endian to an uint64
// value.
func getUint64LE(b []byte) uint64 {
	x := uint64(b[7]) << 56
	x |= uint64(b[6]) << 48
	x |= uint64(b[5]) << 40
	x |= uint64(b[4]) << 32
	x |= uint64(b[3]) << 24
	x |= uint64(b[2]) << 16
	x |= uint64(b[1]) << 8
	x |= uint64(b[0])
	return x
}

// readUint64LE reads a uint64 little-endian integer from reader.
func readUint64LE(r io.Reader) (x uint64, err error) {
	b := make([]byte, 8)
	if _, err = io.ReadFull(r, b); err != nil {
		return 0, err
	}
	x = getUint64LE(b)
	return x, nil
}

// initProbSlice initializes a slice of probabilities.
func initProbSlice(p []prob) {
	for i := range p {
		p[i] = probInit
	}
}

// Reads reads data from the decoder stream.
//
// The method might block and is not reentrant.
//
// The end of the LZMA stream is indicated by EOF. There might be other errors
// returned. The decoder will not be able to recover from an error returned.
func (d *Decoder) Read(p []byte) (n int, err error) {
	for {
		var k int
		k, err = d.dict.Read(p[n:])
		n += k
		switch {
		case err == io.EOF:
			if n <= 0 {
				return 0, io.EOF
			}
			return n, nil
		case err != nil:
			return n, err
		case n == len(p):
			return n, nil
		}
		if err = d.fill(); err != nil {
			return n, err
		}
	}
}

// errUnexpectedEOS indicates that the function decoded an unexpected end of
// stream marker
var errUnexpectedEOS = errors.New("unexpected end of stream marker")

// fill puts at lest the requested number of bytes into the decoder dictionary.
func (d *Decoder) fill() error {
	for !d.dict.eof && d.dict.readable() < d.dict.b {
		op, err := d.decodeOp()
		if err != nil {
			switch {
			case err == eofDecoded:
				if d.unpackLen != noUnpackLen &&
					d.decodedLen != d.unpackLen {
					return errUnexpectedEOS
				}
				d.dict.eof = true
				return nil
			case err == io.EOF:
				return errors.New(
					"unexpected end of compressed stream")
			default:
				return err
			}
		}
		n := d.decodedLen + uint64(op.Len())
		if n < d.decodedLen {
			panic("negative op length or overflow of decodedLen")
		}
		d.decodedLen = n
		if d.decodedLen >= d.unpackLen {
			if d.decodedLen == d.unpackLen {
				d.dict.eof = true
				return nil
			}
			return errors.New(
				"decoded stream longer than provided length")
		}
		if err = op.applyDecoderDict(d.dict); err != nil {
			return err
		}
	}
	return nil
}

// updateStateLiteral updates the state for a literal.
func (d *Decoder) updateStateLiteral() {
	switch {
	case d.state < 4:
		d.state = 0
		return
	case d.state < 10:
		d.state -= 3
		return
	}
	d.state -= 6
	return
}

// updateStateMatch updates the state for a match.
func (d *Decoder) updateStateMatch() {
	if d.state < 7 {
		d.state = 7
		return
	}
	d.state = 10
	return
}

// updateStateRep updates the state for a repetition.
func (d *Decoder) updateStateRep() {
	if d.state < 7 {
		d.state = 8
	}
	d.state = 11
}

// updateStateShortRep updates the state for a short repetition.
func (d *Decoder) updateStateShortRep() {
	if d.state < 7 {
		d.state = 9
	}
	d.state = 11
}

// decodeLiteral decodes a literal.
func (d *Decoder) decodeLiteral() (op operation, err error) {
	prevByte := d.dict.getByte(1)
	lp, lc := uint(d.properties.LP), uint(d.properties.LC)
	litState := ((uint32(d.dict.total) & ((1 << lp) - 1)) << lc) |
		(uint32(prevByte) >> (8 - lc))
	match := d.dict.getByte(int(d.rep[0]) + 1)
	s, err := d.litDecoder.Decode(d.rd, d.state, match, litState)
	if err != nil {
		return nil, err
	}
	return lit{s}, nil
}

// errWrongTermination indicates that a termination symbol has been received,
// but the range decoder could still produces more data
var errWrongTermination = errors.New(
	"range decoder doesn't support termination")

// decodeOp indicates an EOF of the decoded file
var eofDecoded = errors.New("EOF of decoded stream")

// decodeOp decodes an operation. The function returns eofDecoded if there is
// an explicit termination marker.
func (d *Decoder) decodeOp() (op operation, err error) {
	posState := uint32(d.dict.total) & d.posBitMask
	state2 := (d.state << maxPosBits) | posState

	b, err := d.isMatch[state2].Decode(d.rd)
	if err != nil {
		return nil, err
	}
	if b == 0 {
		// literal
		op, err := d.decodeLiteral()
		if err != nil {
			return nil, err
		}
		d.updateStateLiteral()
		return op, nil
	}
	b, err = d.isRep[d.state].Decode(d.rd)
	if err != nil {
		return nil, err
	}
	if b == 0 {
		// simple match
		d.rep[3], d.rep[2], d.rep[1] = d.rep[2], d.rep[1], d.rep[0]
		d.updateStateMatch()
		// The length decoder returns the length offset.
		l, err := d.lengthDecoder.Decode(d.rd, posState)
		if err != nil {
			return nil, err
		}
		// The dist decoder returns the distance offset. The actual
		// distance is 1 higher.
		d.rep[0], err = d.distDecoder.Decode(l, d.rd)
		if err != nil {
			return nil, err
		}
		if d.rep[0] == 0xffffffff {
			if !d.rd.possiblyAtEnd() {
				return nil, errWrongTermination
			}
			return nil, eofDecoded
		}
		op = rep{length: int(l) + minLength,
			distance: int(d.rep[0]) + minDistance}
		return op, nil
	}
	b, err = d.isRepG0[d.state].Decode(d.rd)
	if err != nil {
		return nil, err
	}
	dist := d.rep[0]
	if b == 0 {
		// rep match 0
		b, err = d.isRepG0Long[state2].Decode(d.rd)
		if err != nil {
			return nil, err
		}
		if b == 0 {
			d.updateStateShortRep()
			op = rep{length: 1,
				distance: int(d.rep[0]) + minDistance}
			return op, nil
		}
	} else {
		b, err = d.isRepG1[d.state].Decode(d.rd)
		if err != nil {
			return nil, err
		}
		if b == 0 {
			dist = d.rep[1]
		} else {
			b, err = d.isRepG2[d.state].Decode(d.rd)
			if err != nil {
				return nil, err
			}
			if b == 0 {
				dist = d.rep[2]
			} else {
				dist = d.rep[3]
				d.rep[3] = d.rep[2]
			}
			d.rep[2] = d.rep[1]
		}
		d.rep[1] = d.rep[0]
		d.rep[0] = dist
	}
	l, err := d.repLengthDecoder.Decode(d.rd, posState)
	if err != nil {
		return nil, err
	}
	d.updateStateRep()
	op = rep{length: int(l) + minLength, distance: int(dist) + minDistance}
	return op, nil
}
<|endoftext|>


<commit_before>package etcd

import (
	"bytes"
	"errors"
	"os"
	"path"

	"github.com/barakmich/agro"
	"github.com/barakmich/agro/blockset"
	"github.com/barakmich/agro/models"
	"golang.org/x/net/context"

	// TODO(barakmich): And this is why vendoring sucks. I shouldn't need to
	//import this, but I do, because I'm using etcdserverpb from head, and *it*
	//expects its own vendored version. Admittedly, this should get better with
	//GO15VENDORING, but etcd doesn't support that yet.
	"github.com/coreos/etcd/Godeps/_workspace/src/google.golang.org/grpc"

	pb "github.com/coreos/etcd/etcdserver/etcdserverpb"
)

const keyPrefix = "/github.com/barakmich/agro/"

func init() {
	agro.RegisterMetadataService("etcd", newEtcdMetadata)
}

type etcd struct {
	cfg           agro.Config
	global        agro.GlobalMetadata
	volumeprinter agro.VolumeID
	inodeprinter  agro.INodeID
	volumesCache  map[string]agro.VolumeID

	conn *grpc.ClientConn
	kv   pb.KVClient
}

func newEtcdMetadata(cfg agro.Config) (agro.MetadataService, error) {
	conn, err := grpc.Dial(cfg.MetadataAddress)
	if err != nil {
		return nil, err
	}
	client := pb.NewKVClient(conn)

	//TODO(barakmich): Fetch global metadata. If it doesn't exist, we haven't run mkfs yet. For now, we'll just hardcode it.
	global := agro.GlobalMetadata{
		BlockSize:        8 * 1024,
		DefaultBlockSpec: agro.BlockLayerSpec{blockset.CRC, blockset.Base},
	}
	return &etcd{
		cfg:          cfg,
		global:       global,
		conn:         conn,
		kv:           client,
		volumesCache: make(map[string]agro.VolumeID),
	}, nil
}

func (e *etcd) GlobalMetadata() (agro.GlobalMetadata, error) {
	return e.global, nil
}

func (e *etcd) Close() error {
	return e.conn.Close()
}

func (e *etcd) CreateVolume(volume string) error {
	key := agro.Path{Volume: volume, Path: "/"}
	tx := tx().If(
		keyEquals(mkKey("meta", "volumeprinter"), uint64ToBytes(uint64(e.volumeprinter))),
	).Then(
		setKey(mkKey("meta", "volumeprinter"), uint64ToBytes(uint64(e.volumeprinter+1))),
		setKey(mkKey("volumes", volume), uint64ToBytes(uint64(e.volumeprinter+1))),
		setKey(mkKey("dirs", key.Key()), newDirProto(&models.Metadata{})),
	).Else(
		getKey(mkKey("meta", "volumeprinter")),
	).Tx()
	resp, err := e.kv.Txn(context.Background(), tx)
	if err != nil {
		return err
	}
	if !resp.Succeeded {
		e.volumeprinter = agro.VolumeID(bytesToUint64(resp.Responses[0].ResponseRange.Kvs[0].Value))
		return agro.ErrAgain
	}
	e.volumeprinter++
	return nil
}

func (e *etcd) GetVolumes() ([]string, error) {
	req := getPrefix(mkKey("volumes"))
	resp, err := e.kv.Range(context.Background(), req)
	if err != nil {
		return nil, err
	}
	var out []string
	for _, x := range resp.Kvs {
		p := string(x.Key)
		out = append(out, path.Base(p))
	}
	return out, nil
}

func (e *etcd) GetVolumeID(volume string) (agro.VolumeID, error) {
	if v, ok := e.volumesCache[volume]; ok {
		return v, nil
	}
	req := getKey(mkKey("volumes", volume))
	resp, err := e.kv.Range(context.Background(), req)
	if err != nil {
		return 0, err
	}
	if resp.More {
		// What do?
		return 0, errors.New("implement me")
	}
	if len(resp.Kvs) == 0 {
		return 0, errors.New("etcd: no such volume exists")
	}
	vid := agro.VolumeID(bytesToUint64(resp.Kvs[0].Value))
	e.volumesCache[volume] = vid
	return vid, nil

}

func (e *etcd) CommitInodeIndex() (agro.INodeID, error) {
	tx := tx().If(
		keyEquals(mkKey("meta", "inodeprinter"), uint64ToBytes(uint64(e.inodeprinter))),
	).Then(
		setKey(mkKey("meta", "inodeprinter"), uint64ToBytes(uint64(e.inodeprinter+1))),
	).Else(
		getKey(mkKey("meta", "inodeprinter")),
	).Tx()
	resp, err := e.kv.Txn(context.Background(), tx)
	if err != nil {
		return 0, err
	}
	if !resp.Succeeded {
		e.inodeprinter = agro.INodeID(bytesToUint64(resp.Responses[0].ResponseRange.Kvs[0].Value))
		return 0, agro.ErrAgain
	}
	i := e.inodeprinter
	e.inodeprinter++
	return i, nil
}

func (e *etcd) Mkdir(path agro.Path, dir *models.Directory) error {
	super, ok := path.Super()
	if !ok {
		return errors.New("etcd: not a directory")
	}
	tx := tx().If(
		keyExists(mkKey("dirs", super.Key())),
	).Then(
		setKey(mkKey("dirs", path.Key()), newDirProto(&models.Metadata{})),
	).Tx()
	resp, err := e.kv.Txn(context.Background(), tx)
	if err != nil {
		return err
	}
	if !resp.Succeeded {
		return os.ErrNotExist
	}
	return nil
}

func (e *etcd) getDirRaw(p agro.Path) (*pb.TxnResponse, error) {
	tx := tx().If(
		keyExists(mkKey("dirs", p.Key())),
	).Then(
		getKey(mkKey("dirs", p.Key())),
		getPrefix(mkKey("dirs", p.SubdirsPrefix())),
	).Tx()
	return e.kv.Txn(context.Background(), tx)
}

func (e *etcd) Getdir(p agro.Path) (*models.Directory, []agro.Path, error) {
	resp, err := e.getDirRaw(p)
	if err != nil {
		return nil, nil, err
	}
	if !resp.Succeeded {
		return nil, nil, os.ErrNotExist
	}
	dirkv := resp.Responses[0].ResponseRange.Kvs[0]
	outdir := &models.Directory{}
	err = outdir.Unmarshal(dirkv.Value)
	if err != nil {
		return nil, nil, err
	}
	var outpaths []agro.Path
	for _, kv := range resp.Responses[1].ResponseRange.Kvs {
		s := bytes.SplitN(kv.Key, []byte{':'}, 2)
		outpaths = append(outpaths, agro.Path{
			Volume: p.Volume,
			Path:   string(s[2]),
		})
	}
	return outdir, outpaths, nil
}

func (e *etcd) SetFileINode(p agro.Path, ref agro.INodeRef) error {
	resp, err := e.getDirRaw(p)
	if err != nil {
		return err
	}
	if !resp.Succeeded {
		panic("shouldn't be able to SetFileINode a non-existent dir")
	}
	return e.trySetFileINode(p, ref, resp)
}

func (e *etcd) trySetFileINode(p agro.Path, ref agro.INodeRef, resp *pb.TxnResponse) error {
	dirkv := resp.Responses[0].ResponseRange.Kvs[0]
	dir := &models.Directory{}
	err := dir.Unmarshal(dirkv.Value)
	if err != nil {
		return err
	}
	if dir.Files == nil {
		dir.Files = make(map[string]uint64)
	}
	dir.Files[p.Filename()] = uint64(ref.INode)
	b, err := dir.Marshal()
	tx := tx().If(
		keyIsVersion(dirkv.Key, dirkv.Version),
	).Then(
		setKey(dirkv.Key, b),
	).Else(
		getKey(dirkv.Key),
	).Tx()
	resp, err = e.kv.Txn(context.Background(), tx)
	if err != nil {
		return err
	}
	if resp.Succeeded {
		return nil
	}
	return e.trySetFileINode(p, ref, resp)
}

func (e *etcd) Mkfs(gmd agro.GlobalMetadata) error {
	tx := tx().If(
		keyNotExists(mkKey("meta", "volumeprinter")),
	).Then(
		setKey(mkKey("meta", "volumeprinter"), uint64ToBytes(1)),
		setKey(mkKey("meta", "inodeprinter"), uint64ToBytes(1)),
	).Else(
		getKey(mkKey("meta", "volumeprinter")),
		getKey(mkKey("meta", "inodeprinter")),
	).Tx()
	resp, err := e.kv.Txn(context.Background(), tx)
	if err != nil {
		return err
	}
	if !resp.Succeeded {
		e.volumeprinter = agro.VolumeID(bytesToUint64(resp.Responses[0].ResponseRange.Kvs[0].Value))
		e.inodeprinter = agro.INodeID(bytesToUint64(resp.Responses[1].ResponseRange.Kvs[0].Value))
		return nil
	}
	e.volumeprinter = 1
	e.inodeprinter = 1
	return nil
}
<commit_msg>and implement etcd extensions<commit_after>package etcd

import (
	"bytes"
	"errors"
	"os"
	"path"
	"time"

	"github.com/barakmich/agro"
	"github.com/barakmich/agro/blockset"
	"github.com/barakmich/agro/metadata"
	"github.com/barakmich/agro/models"
	"golang.org/x/net/context"

	// TODO(barakmich): And this is why vendoring sucks. I shouldn't need to
	//import this, but I do, because I'm using etcdserverpb from head, and *it*
	//expects its own vendored version. Admittedly, this should get better with
	//GO15VENDORING, but etcd doesn't support that yet.
	"github.com/coreos/etcd/Godeps/_workspace/src/google.golang.org/grpc"
	"github.com/coreos/pkg/capnslog"

	pb "github.com/coreos/etcd/etcdserver/etcdserverpb"
)

var clog = capnslog.NewPackageLogger("github.com/barakmich/agro", "etcd")

const (
	keyPrefix      = "/github.com/barakmich/agro/"
	peerTimeoutMax = 5 * time.Second
)

func init() {
	agro.RegisterMetadataService("etcd", newEtcdMetadata)
}

type etcd struct {
	cfg           agro.Config
	global        agro.GlobalMetadata
	volumeprinter agro.VolumeID
	inodeprinter  agro.INodeID
	volumesCache  map[string]agro.VolumeID

	conn *grpc.ClientConn
	kv   pb.KVClient
	uuid string
}

func newEtcdMetadata(cfg agro.Config) (agro.MetadataService, error) {
	uuid, err := metadata.MakeOrGetUUID(cfg.DataDir)
	if err != nil {
		return nil, err
	}

	conn, err := grpc.Dial(cfg.MetadataAddress)
	if err != nil {
		return nil, err
	}
	client := pb.NewKVClient(conn)

	//TODO(barakmich): Fetch global metadata. If it doesn't exist, we haven't run mkfs yet. For now, we'll just hardcode it.
	global := agro.GlobalMetadata{
		BlockSize:        8 * 1024,
		DefaultBlockSpec: agro.BlockLayerSpec{blockset.CRC, blockset.Base},
	}
	return &etcd{
		cfg:          cfg,
		global:       global,
		conn:         conn,
		kv:           client,
		volumesCache: make(map[string]agro.VolumeID),
		uuid:         uuid,
	}, nil
}

func (e *etcd) GlobalMetadata() (agro.GlobalMetadata, error) {
	return e.global, nil
}

func (e *etcd) UUID() string {
	return e.uuid
}

func (e *etcd) RegisterPeer(p *models.PeerInfo) error {
	p.LastSeen = time.Now().UnixNano()
	data, err := p.Marshal()
	if err != nil {
		return err
	}
	_, err = e.kv.Put(context.Background(),
		setKey(mkKey("nodes", p.UUID), data),
	)
	return err
}

func (e *etcd) GetPeers() ([]*models.PeerInfo, error) {
	resp, err := e.kv.Range(context.Background(), getPrefix(mkKey("nodes")))
	if err != nil {
		return nil, err
	}
	var out []*models.PeerInfo
	for _, x := range resp.Kvs {
		var p models.PeerInfo
		err := p.Unmarshal(x.Value)
		if err != nil {
			// Intentionally ignore a peer that doesn't unmarshal properly.
			clog.Errorf("peer at key %s didn't unmarshal correctly", string(x.Key))
			continue
		}
		if time.Since(time.Unix(0, p.LastSeen)) > peerTimeoutMax {
			clog.Debugf("peer at key %s didn't unregister; fixed with leases in etcdv3", string(x.Key))
			continue
		}
		out = append(out, &p)
	}
	return out, nil
}

func (e *etcd) Close() error {
	return e.conn.Close()
}

func (e *etcd) CreateVolume(volume string) error {
	key := agro.Path{Volume: volume, Path: "/"}
	tx := tx().If(
		keyEquals(mkKey("meta", "volumeprinter"), uint64ToBytes(uint64(e.volumeprinter))),
	).Then(
		setKey(mkKey("meta", "volumeprinter"), uint64ToBytes(uint64(e.volumeprinter+1))),
		setKey(mkKey("volumes", volume), uint64ToBytes(uint64(e.volumeprinter+1))),
		setKey(mkKey("dirs", key.Key()), newDirProto(&models.Metadata{})),
	).Else(
		getKey(mkKey("meta", "volumeprinter")),
	).Tx()
	resp, err := e.kv.Txn(context.Background(), tx)
	if err != nil {
		return err
	}
	if !resp.Succeeded {
		e.volumeprinter = agro.VolumeID(bytesToUint64(resp.Responses[0].ResponseRange.Kvs[0].Value))
		return agro.ErrAgain
	}
	e.volumeprinter++
	return nil
}

func (e *etcd) GetVolumes() ([]string, error) {
	req := getPrefix(mkKey("volumes"))
	resp, err := e.kv.Range(context.Background(), req)
	if err != nil {
		return nil, err
	}
	var out []string
	for _, x := range resp.Kvs {
		p := string(x.Key)
		out = append(out, path.Base(p))
	}
	return out, nil
}

func (e *etcd) GetVolumeID(volume string) (agro.VolumeID, error) {
	if v, ok := e.volumesCache[volume]; ok {
		return v, nil
	}
	req := getKey(mkKey("volumes", volume))
	resp, err := e.kv.Range(context.Background(), req)
	if err != nil {
		return 0, err
	}
	if resp.More {
		// What do?
		return 0, errors.New("implement me")
	}
	if len(resp.Kvs) == 0 {
		return 0, errors.New("etcd: no such volume exists")
	}
	vid := agro.VolumeID(bytesToUint64(resp.Kvs[0].Value))
	e.volumesCache[volume] = vid
	return vid, nil

}

func (e *etcd) CommitInodeIndex() (agro.INodeID, error) {
	tx := tx().If(
		keyEquals(mkKey("meta", "inodeprinter"), uint64ToBytes(uint64(e.inodeprinter))),
	).Then(
		setKey(mkKey("meta", "inodeprinter"), uint64ToBytes(uint64(e.inodeprinter+1))),
	).Else(
		getKey(mkKey("meta", "inodeprinter")),
	).Tx()
	resp, err := e.kv.Txn(context.Background(), tx)
	if err != nil {
		return 0, err
	}
	if !resp.Succeeded {
		e.inodeprinter = agro.INodeID(bytesToUint64(resp.Responses[0].ResponseRange.Kvs[0].Value))
		return 0, agro.ErrAgain
	}
	i := e.inodeprinter
	e.inodeprinter++
	return i, nil
}

func (e *etcd) Mkdir(path agro.Path, dir *models.Directory) error {
	super, ok := path.Super()
	if !ok {
		return errors.New("etcd: not a directory")
	}
	tx := tx().If(
		keyExists(mkKey("dirs", super.Key())),
	).Then(
		setKey(mkKey("dirs", path.Key()), newDirProto(&models.Metadata{})),
	).Tx()
	resp, err := e.kv.Txn(context.Background(), tx)
	if err != nil {
		return err
	}
	if !resp.Succeeded {
		return os.ErrNotExist
	}
	return nil
}

func (e *etcd) getDirRaw(p agro.Path) (*pb.TxnResponse, error) {
	tx := tx().If(
		keyExists(mkKey("dirs", p.Key())),
	).Then(
		getKey(mkKey("dirs", p.Key())),
		getPrefix(mkKey("dirs", p.SubdirsPrefix())),
	).Tx()
	return e.kv.Txn(context.Background(), tx)
}

func (e *etcd) Getdir(p agro.Path) (*models.Directory, []agro.Path, error) {
	resp, err := e.getDirRaw(p)
	if err != nil {
		return nil, nil, err
	}
	if !resp.Succeeded {
		return nil, nil, os.ErrNotExist
	}
	dirkv := resp.Responses[0].ResponseRange.Kvs[0]
	outdir := &models.Directory{}
	err = outdir.Unmarshal(dirkv.Value)
	if err != nil {
		return nil, nil, err
	}
	var outpaths []agro.Path
	for _, kv := range resp.Responses[1].ResponseRange.Kvs {
		s := bytes.SplitN(kv.Key, []byte{':'}, 2)
		outpaths = append(outpaths, agro.Path{
			Volume: p.Volume,
			Path:   string(s[2]),
		})
	}
	return outdir, outpaths, nil
}

func (e *etcd) SetFileINode(p agro.Path, ref agro.INodeRef) error {
	resp, err := e.getDirRaw(p)
	if err != nil {
		return err
	}
	if !resp.Succeeded {
		panic("shouldn't be able to SetFileINode a non-existent dir")
	}
	return e.trySetFileINode(p, ref, resp)
}

func (e *etcd) trySetFileINode(p agro.Path, ref agro.INodeRef, resp *pb.TxnResponse) error {
	dirkv := resp.Responses[0].ResponseRange.Kvs[0]
	dir := &models.Directory{}
	err := dir.Unmarshal(dirkv.Value)
	if err != nil {
		return err
	}
	if dir.Files == nil {
		dir.Files = make(map[string]uint64)
	}
	dir.Files[p.Filename()] = uint64(ref.INode)
	b, err := dir.Marshal()
	tx := tx().If(
		keyIsVersion(dirkv.Key, dirkv.Version),
	).Then(
		setKey(dirkv.Key, b),
	).Else(
		getKey(dirkv.Key),
	).Tx()
	resp, err = e.kv.Txn(context.Background(), tx)
	if err != nil {
		return err
	}
	if resp.Succeeded {
		return nil
	}
	return e.trySetFileINode(p, ref, resp)
}

func (e *etcd) Mkfs(gmd agro.GlobalMetadata) error {
	tx := tx().If(
		keyNotExists(mkKey("meta", "volumeprinter")),
	).Then(
		setKey(mkKey("meta", "volumeprinter"), uint64ToBytes(1)),
		setKey(mkKey("meta", "inodeprinter"), uint64ToBytes(1)),
	).Else(
		getKey(mkKey("meta", "volumeprinter")),
		getKey(mkKey("meta", "inodeprinter")),
	).Tx()
	resp, err := e.kv.Txn(context.Background(), tx)
	if err != nil {
		return err
	}
	if !resp.Succeeded {
		e.volumeprinter = agro.VolumeID(bytesToUint64(resp.Responses[0].ResponseRange.Kvs[0].Value))
		e.inodeprinter = agro.INodeID(bytesToUint64(resp.Responses[1].ResponseRange.Kvs[0].Value))
		return nil
	}
	e.volumeprinter = 1
	e.inodeprinter = 1
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/cenkalti/backoff"
	"github.com/fsouza/go-dockerclient"
)

var (
	lastStartContainerCall = time.Now()
)

func getImages(dangling bool) []docker.APIImages {
	images, err := dockerClient.ListImages(docker.ListImagesOptions{
		All: false,
	})
	orFail(err)

	nowSeconds := time.Now().Unix()
	response := []docker.APIImages{}

	for _, v := range images {
		if dangling {
			if len(v.RepoTags) == 1 && v.RepoTags[0] == "<none>:<none>" && (nowSeconds-v.Created) > 3600 {
				response = append(response, v)
			}
		} else {
			if len(v.RepoTags) != 1 || v.RepoTags[0] != "<none>:<none>" {
				response = append(response, v)
			}
		}
	}

	return response
}

func cleanDangling() {
	images := getImages(true)
	for _, v := range images {
		log.Printf("Removing dangling image: %s", v.ID)
		dockerClient.RemoveImage(v.ID)
	}
}

func removeNotRequiredImages() {
	required := []string{}
	for _, v := range *cfg {
		required = append(required, fmt.Sprintf("%s:%s", v.Image, v.Tag))
	}

	currentImages := getImages(false)
	for _, i := range currentImages {
		found := false
		for _, t := range i.RepoTags {
			for _, r := range required {
				if r == t {
					found = true
				}
			}
		}

		if !found {
			log.Printf("Removing not required image: %s", i.ID)
			dockerClient.RemoveImage(i.ID)
		}
	}
}

func refreshImages() {
	for _, v := range *cfg {
		pullImage(v.Image, v.Tag)
	}
}

func pullImage(image, tag string) {
	auth := docker.AuthConfiguration{}

	reginfo := strings.SplitN(image, "/", 2)
	if len(reginfo) == 2 {
		for s, a := range authConfig.Configs {
			if strings.Contains(s, fmt.Sprintf("://%s/", reginfo[0])) {
				auth = a
			}
		}
	}

	log.Printf("Refreshing repo %s:%s...", image, tag)
	err := dockerClient.PullImage(docker.PullImageOptions{
		Repository: image,
		Tag:        tag,
	}, auth)
	orLog(err)
}

func bootContainer(name string, cfg containerConfig) {
	var (
		container *docker.Container
		err       error
		bo        = backoff.NewExponentialBackOff()
	)

	cs, err := cfg.checksum()
	orFail(err)
	if err != nil {
		return
	}

	labels := cfg.Labels
	if labels == nil {
		labels = map[string]string{}
	}
	labels["io.luzifer.dockermanager.cfghash"] = cs
	labels["io.luzifer.dockermanager.managed"] = "true"

	if cfg.StartTimes != "" {
		labels["io.luzifer.dockermanager.scheduler"] = "true"
	}

	newcfg := &docker.Config{
		AttachStdin:  false,
		AttachStdout: true,
		AttachStderr: true,
		Image:        strings.Join([]string{cfg.Image, cfg.Tag}, ":"),
		Env:          cfg.Environment,
		Cmd:          cfg.Command,
		Labels:       labels,
	}

	bo.MaxElapsedTime = time.Minute
	err = backoff.Retry(func() error {
		log.Printf("Creating container %s", name)
		container, err = dockerClient.CreateContainer(docker.CreateContainerOptions{
			Name:   name,
			Config: newcfg,
		})
		orLog(err)
		if err != nil {
			if strings.Contains(err.Error(), "no such image") {
				pullImage(cfg.Image, cfg.Tag)
			}
			if strings.Contains(err.Error(), "is already in use by container") {
				cleanContainers()
			}
		}
		return err
	}, bo)

	if err != nil {
		log.Printf("Unable to start container '%s': %s", name, err)
		return
	}

	hostConfig := docker.HostConfig{
		Binds:        cfg.Volumes,
		Links:        cfg.Links,
		Privileged:   false,
		PortBindings: make(map[docker.Port][]docker.PortBinding),
	}

	for _, v := range cfg.Ports {
		s := strings.Split(v.Local, ":")
		hostConfig.PortBindings[docker.Port(v.Container)] = []docker.PortBinding{docker.PortBinding{
			HostIP:   s[0],
			HostPort: s[1],
		}}
	}

	log.Printf("Starting container %s", container.Name)
	dockerClient.StartContainer(container.Name, &hostConfig)
}

func listRunningContainers() ([]docker.APIContainers, error) {
	containers, err := dockerClient.ListContainers(docker.ListContainersOptions{
		All: false,
	})
	orLog(err)
	return containers, err
}

func getExpectedRunningNames() []string {
	expectedRunning := []string{}
	for name, containerCfg := range *cfg {
		if containerCfg.shouldBeRunning(serfElector.MyName) {
			expectedRunning = append(expectedRunning, name)
		}
	}
	return expectedRunning
}

func stopUnexpectedContainers() {
	expectedRunning := getExpectedRunningNames()

	// Stop containers not expected to be running
	currentRunning, err := listRunningContainers()
	orLog(err)
	if err != nil {
		return
	}
	for _, v := range currentRunning {
		allowed := false

		containerDetails, err := dockerClient.InspectContainer(v.ID)
		orFail(err)

		_, isManaged := containerDetails.Config.Labels["io.luzifer.dockermanager.managed"]
		if !params.ManageFullHost && !isManaged {
			allowed = true
		}

		if containerDetails.Config.Labels["io.luzifer.dockermanager.scheduler"] == "true" {
			// If it's a scheduled container keep it running
			allowed = true
		}

		for _, n := range v.Names {
			if stringInSlice(strings.Trim(n, "/"), expectedRunning) {
				allowed = true
			}
		}
		if !allowed {
			log.Printf("Stopping container %s as it is not expected to be running.", v.ID)
			err := dockerClient.StopContainer(v.ID, 5)
			orFail(err)
			time.Sleep(time.Second * 5)
		}
	}
}

func removeDeprecatedContainers() {
	expectedRunning := getExpectedRunningNames()

	// Stop containers who are of deprecated images
	currentRunning, err := listRunningContainers()
	orLog(err)
	if err != nil {
		return
	}
	images := getImages(false)
	for _, n := range expectedRunning {
		repoName := strings.Join([]string{(*cfg)[n].Image, (*cfg)[n].Tag}, ":")
		currentImageID := "0"
		for _, i := range images {
			if stringInSlice(repoName, i.RepoTags) {
				currentImageID = i.ID
			}
		}
		if currentImageID == "0" {
			log.Printf("Found no image ID for repo %s", strings.Join([]string{(*cfg)[n].Image, (*cfg)[n].Tag}, ":"))
			continue
		}
		for _, v := range currentRunning {
			containerDetails, err := dockerClient.InspectContainer(v.ID)
			orFail(err)

			cs, err := (*cfg)[n].checksum()

			if stringInSlice(fmt.Sprintf("/%s", n), v.Names) {
				needsUpdate := false
				if !strings.HasPrefix(currentImageID, containerDetails.Image) {
					log.Printf("Container %s has a new image version.", n)
					needsUpdate = true
				}
				if containerDetails.Config.Labels["io.luzifer.dockermanager.cfghash"] != cs {
					log.Printf("Container %s has a configuration update.", n)
					needsUpdate = true
				}
				if needsUpdate && !timeAllowed((*cfg)[n].UpdateTimes) {
					log.Printf("Image %s has update but container %s (%s) is not allowed to update now.", v.Image, n, v.ID)
					needsUpdate = false
				}

				if needsUpdate {
					stopWaitTime := (*cfg)[n].StopTimeout
					if stopWaitTime == 0 {
						stopWaitTime = 5
					}

					log.Printf("Image: %s Current: %s", containerDetails.Image, currentImageID)
					log.Printf("Stopping deprecated container %s", v.ID)
					err := dockerClient.StopContainer(v.ID, stopWaitTime)
					orFail(err)

					log.Printf("Removing deprecated container %s", v.ID)
					dockerClient.RemoveContainer(docker.RemoveContainerOptions{
						ID: v.ID,
					})

					currentRunning, err = listRunningContainers()
					orFail(err)
				}
			}
		}
	}
}

func startExpectedContainers() {
	expectedRunning := getExpectedRunningNames()
	lastStartContainerCall = time.Now()

	// Start expected containers
	currentRunning, err := listRunningContainers()
	orLog(err)
	if err != nil {
		return
	}
	runningNames := []string{}
	for _, v := range currentRunning {
		for _, n := range v.Names {
			runningNames = append(runningNames, strings.Trim(n, "/"))
		}
	}
	for _, n := range expectedRunning {
		if !stringInSlice(n, runningNames) {
			bootContainer(n, (*cfg)[n])
		}
	}
}

func cleanContainers() {
	runningContainers, err := dockerClient.ListContainers(docker.ListContainersOptions{
		All: true,
	})
	orLog(err)
	if err != nil {
		return
	}

	for _, v := range runningContainers {
		if strings.HasPrefix(v.Status, "Exited") {
			log.Printf("Removing stopped container %s", v.ID)
			err := dockerClient.RemoveContainer(docker.RemoveContainerOptions{
				ID: v.ID,
			})
			orLog(err)
		}
	}
}
<commit_msg>Do not keep images not required for current host<commit_after>package main

import (
	"fmt"
	"log"
	"strings"
	"time"

	"github.com/cenkalti/backoff"
	"github.com/fsouza/go-dockerclient"
)

var (
	lastStartContainerCall = time.Now()
)

func getImages(dangling bool) []docker.APIImages {
	images, err := dockerClient.ListImages(docker.ListImagesOptions{
		All: false,
	})
	orFail(err)

	nowSeconds := time.Now().Unix()
	response := []docker.APIImages{}

	for _, v := range images {
		if dangling {
			if len(v.RepoTags) == 1 && v.RepoTags[0] == "<none>:<none>" && (nowSeconds-v.Created) > 3600 {
				response = append(response, v)
			}
		} else {
			if len(v.RepoTags) != 1 || v.RepoTags[0] != "<none>:<none>" {
				response = append(response, v)
			}
		}
	}

	return response
}

func cleanDangling() {
	images := getImages(true)
	for _, v := range images {
		log.Printf("Removing dangling image: %s", v.ID)
		dockerClient.RemoveImage(v.ID)
	}
}

func removeNotRequiredImages() {
	required := []string{}
	for _, v := range *cfg {
		if !stringInSlice(serfElector.MyName, v.Hosts) {
			continue
		}
		required = append(required, fmt.Sprintf("%s:%s", v.Image, v.Tag))
	}

	currentImages := getImages(false)
	for _, i := range currentImages {
		found := false
		for _, t := range i.RepoTags {
			for _, r := range required {
				if r == t {
					found = true
				}
			}
		}

		if !found {
			log.Printf("Removing not required image: %s", i.ID)
			dockerClient.RemoveImage(i.ID)
		}
	}
}

func refreshImages() {
	for _, v := range *cfg {
		if !stringInSlice(serfElector.MyName, v.Hosts) {
			continue
		}
		pullImage(v.Image, v.Tag)
	}
}

func pullImage(image, tag string) {
	auth := docker.AuthConfiguration{}

	reginfo := strings.SplitN(image, "/", 2)
	if len(reginfo) == 2 {
		for s, a := range authConfig.Configs {
			if strings.Contains(s, fmt.Sprintf("://%s/", reginfo[0])) {
				auth = a
			}
		}
	}

	log.Printf("Refreshing repo %s:%s...", image, tag)
	err := dockerClient.PullImage(docker.PullImageOptions{
		Repository: image,
		Tag:        tag,
	}, auth)
	orLog(err)
}

func bootContainer(name string, cfg containerConfig) {
	var (
		container *docker.Container
		err       error
		bo        = backoff.NewExponentialBackOff()
	)

	cs, err := cfg.checksum()
	orFail(err)
	if err != nil {
		return
	}

	labels := cfg.Labels
	if labels == nil {
		labels = map[string]string{}
	}
	labels["io.luzifer.dockermanager.cfghash"] = cs
	labels["io.luzifer.dockermanager.managed"] = "true"

	if cfg.StartTimes != "" {
		labels["io.luzifer.dockermanager.scheduler"] = "true"
	}

	newcfg := &docker.Config{
		AttachStdin:  false,
		AttachStdout: true,
		AttachStderr: true,
		Image:        strings.Join([]string{cfg.Image, cfg.Tag}, ":"),
		Env:          cfg.Environment,
		Cmd:          cfg.Command,
		Labels:       labels,
	}

	bo.MaxElapsedTime = time.Minute
	err = backoff.Retry(func() error {
		log.Printf("Creating container %s", name)
		container, err = dockerClient.CreateContainer(docker.CreateContainerOptions{
			Name:   name,
			Config: newcfg,
		})
		orLog(err)
		if err != nil {
			if strings.Contains(err.Error(), "no such image") {
				pullImage(cfg.Image, cfg.Tag)
			}
			if strings.Contains(err.Error(), "is already in use by container") {
				cleanContainers()
			}
		}
		return err
	}, bo)

	if err != nil {
		log.Printf("Unable to start container '%s': %s", name, err)
		return
	}

	hostConfig := docker.HostConfig{
		Binds:        cfg.Volumes,
		Links:        cfg.Links,
		Privileged:   false,
		PortBindings: make(map[docker.Port][]docker.PortBinding),
	}

	for _, v := range cfg.Ports {
		s := strings.Split(v.Local, ":")
		hostConfig.PortBindings[docker.Port(v.Container)] = []docker.PortBinding{docker.PortBinding{
			HostIP:   s[0],
			HostPort: s[1],
		}}
	}

	log.Printf("Starting container %s", container.Name)
	dockerClient.StartContainer(container.Name, &hostConfig)
}

func listRunningContainers() ([]docker.APIContainers, error) {
	containers, err := dockerClient.ListContainers(docker.ListContainersOptions{
		All: false,
	})
	orLog(err)
	return containers, err
}

func getExpectedRunningNames() []string {
	expectedRunning := []string{}
	for name, containerCfg := range *cfg {
		if containerCfg.shouldBeRunning(serfElector.MyName) {
			expectedRunning = append(expectedRunning, name)
		}
	}
	return expectedRunning
}

func stopUnexpectedContainers() {
	expectedRunning := getExpectedRunningNames()

	// Stop containers not expected to be running
	currentRunning, err := listRunningContainers()
	orLog(err)
	if err != nil {
		return
	}
	for _, v := range currentRunning {
		allowed := false

		containerDetails, err := dockerClient.InspectContainer(v.ID)
		orFail(err)

		_, isManaged := containerDetails.Config.Labels["io.luzifer.dockermanager.managed"]
		if !params.ManageFullHost && !isManaged {
			allowed = true
		}

		if containerDetails.Config.Labels["io.luzifer.dockermanager.scheduler"] == "true" {
			// If it's a scheduled container keep it running
			allowed = true
		}

		for _, n := range v.Names {
			if stringInSlice(strings.Trim(n, "/"), expectedRunning) {
				allowed = true
			}
		}
		if !allowed {
			log.Printf("Stopping container %s as it is not expected to be running.", v.ID)
			err := dockerClient.StopContainer(v.ID, 5)
			orFail(err)
			time.Sleep(time.Second * 5)
		}
	}
}

func removeDeprecatedContainers() {
	expectedRunning := getExpectedRunningNames()

	// Stop containers who are of deprecated images
	currentRunning, err := listRunningContainers()
	orLog(err)
	if err != nil {
		return
	}
	images := getImages(false)
	for _, n := range expectedRunning {
		repoName := strings.Join([]string{(*cfg)[n].Image, (*cfg)[n].Tag}, ":")
		currentImageID := "0"
		for _, i := range images {
			if stringInSlice(repoName, i.RepoTags) {
				currentImageID = i.ID
			}
		}
		if currentImageID == "0" {
			log.Printf("Found no image ID for repo %s", strings.Join([]string{(*cfg)[n].Image, (*cfg)[n].Tag}, ":"))
			continue
		}
		for _, v := range currentRunning {
			containerDetails, err := dockerClient.InspectContainer(v.ID)
			orFail(err)

			cs, err := (*cfg)[n].checksum()

			if stringInSlice(fmt.Sprintf("/%s", n), v.Names) {
				needsUpdate := false
				if !strings.HasPrefix(currentImageID, containerDetails.Image) {
					log.Printf("Container %s has a new image version.", n)
					needsUpdate = true
				}
				if containerDetails.Config.Labels["io.luzifer.dockermanager.cfghash"] != cs {
					log.Printf("Container %s has a configuration update.", n)
					needsUpdate = true
				}
				if needsUpdate && !timeAllowed((*cfg)[n].UpdateTimes) {
					log.Printf("Image %s has update but container %s (%s) is not allowed to update now.", v.Image, n, v.ID)
					needsUpdate = false
				}

				if needsUpdate {
					stopWaitTime := (*cfg)[n].StopTimeout
					if stopWaitTime == 0 {
						stopWaitTime = 5
					}

					log.Printf("Image: %s Current: %s", containerDetails.Image, currentImageID)
					log.Printf("Stopping deprecated container %s", v.ID)
					err := dockerClient.StopContainer(v.ID, stopWaitTime)
					orFail(err)

					log.Printf("Removing deprecated container %s", v.ID)
					dockerClient.RemoveContainer(docker.RemoveContainerOptions{
						ID: v.ID,
					})

					currentRunning, err = listRunningContainers()
					orFail(err)
				}
			}
		}
	}
}

func startExpectedContainers() {
	expectedRunning := getExpectedRunningNames()
	lastStartContainerCall = time.Now()

	// Start expected containers
	currentRunning, err := listRunningContainers()
	orLog(err)
	if err != nil {
		return
	}
	runningNames := []string{}
	for _, v := range currentRunning {
		for _, n := range v.Names {
			runningNames = append(runningNames, strings.Trim(n, "/"))
		}
	}
	for _, n := range expectedRunning {
		if !stringInSlice(n, runningNames) {
			bootContainer(n, (*cfg)[n])
		}
	}
}

func cleanContainers() {
	runningContainers, err := dockerClient.ListContainers(docker.ListContainersOptions{
		All: true,
	})
	orLog(err)
	if err != nil {
		return
	}

	for _, v := range runningContainers {
		if strings.HasPrefix(v.Status, "Exited") {
			log.Printf("Removing stopped container %s", v.ID)
			err := dockerClient.RemoveContainer(docker.RemoveContainerOptions{
				ID: v.ID,
			})
			orLog(err)
		}
	}
}
<|endoftext|>


<commit_before>package hstspreload

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"golang.org/x/net/publicsuffix"
	"net"
	"net/http"
	"net/url"
	"os"
	"strings"
	"syscall"
	"time"
)

const (
	// dialTimeout specifies the amount of time that TCP or TLS connections
	// can take to complete.
	dialTimeout = 10 * time.Second
)

// dialer is a global net.Dialer that's used whenever making TLS connections in
// order to enforce dialTimeout.
var dialer = net.Dialer{
	Timeout: dialTimeout,
}

// CheckDomain checks whether the domain passes HSTS preload
// requirements for Chromium. This includes:
//
// - Serving a single HSTS header that passes header requirements.
//
// - Using TLS settings that will not cause new problems for
// Chromium/Chrome users. (Example of a new problem: a missing intermediate certificate
// will turn an error page from overrideable to non-overridable on
// some mobile devices.)
//
// To interpret the result, see the list of conventions in the
// documentation for Issues.
func CheckDomain(domain string) (issues Issues) {
	// Check domain format issues first, since we can report something
	// useful even if the other checks fail.
	issues = combineIssues(issues, checkDomainFormat(domain))

	// We don't currently allow automatic submissions of subdomains.
	eTLD1Issues := checkEffectiveTLDPlusOne(domain)
	issues = combineIssues(issues, eTLD1Issues)

	// Start with an initial probe, and don't do the follow-up checks if
	// we can't connect.
	resp, respIssues := getResponse(domain)
	issues = combineIssues(issues, respIssues)
	if len(respIssues.Errors) == 0 {
		issues = combineIssues(issues, checkChain(certChain(*resp.TLS), domain))
		issues = combineIssues(issues, CheckResponse(*resp))

		// Skip the WWW check if the domain is not eTLD+1.
		if len(eTLD1Issues.Errors) == 0 {
			issues = combineIssues(issues, checkWWW(domain))
		}
	}

	return issues
}

func getResponse(domain string) (resp *http.Response, issues Issues) {
	redirectPrevented := errors.New("REDIRECT_PREVENTED")

	client := &http.Client{
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return redirectPrevented
		},
	}

	resp, err := client.Get("https://" + domain)
	if err != nil {
		if urlError, ok := err.(*url.Error); !ok || urlError.Err != redirectPrevented {
			return resp, issues.addErrorf(
				"Cannot connect using TLS (%q). This might be caused by an incomplete "+
					"certificate chain, which causes issues on mobile devices. "+
					"Check out your site at https://www.ssllabs.com/ssltest/",
				err,
			)
		}
	}

	return resp, issues
}

func checkDomainFormat(domain string) (issues Issues) {
	if strings.HasPrefix(domain, ".") {
		return issues.addErrorf("Domain name error: begins with `.`")
	}
	if strings.HasSuffix(domain, ".") {
		return issues.addErrorf("Domain name error: ends with `.`")
	}
	if strings.Index(domain, "..") != -1 {
		return issues.addErrorf("Domain name error: contains `..`")
	}
	if strings.Count(domain, ".") < 1 {
		return issues.addErrorf("Domain name error: must have at least two labels.")
	}

	domain = strings.ToLower(domain)
	for _, r := range domain {
		if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '-' || r == '.' {
			continue
		}

		return issues.addErrorf("Domain name error: contains invalid characters.")
	}

	return issues
}

func checkEffectiveTLDPlusOne(domain string) (issues Issues) {
	canon, err := publicsuffix.EffectiveTLDPlusOne(domain)
	if err != nil {
		return issues.addErrorf("Internal error: could not compute eTLD+1.")
	}
	if canon != domain {
		return issues.addErrorf(
			"Domain error: `%s` is not eTLD+1. Please preload `%s` instead.",
			domain,
			canon,
		)
	}

	return issues
}

// Takes the domain as argument because we may need to make more network
// connections to see if an ECDSA cert is permissible.
func checkChain(chain []*x509.Certificate, domain string) (issues Issues) {
	issues = combineIssues(issues, checkSHA1(chain))

	return issues
}

func checkSHA1(chain []*x509.Certificate) (issues Issues) {
	if firstSHA1, found := findPropertyInChain(isSHA1, chain); found {
		issues = issues.addErrorf(
			"One or more of the certificates in your certificate chain is signed with SHA-1. "+
				"This needs to be replaced. "+
				"See https://security.googleblog.com/2015/12/an-update-on-sha-1-certificates-in.html. "+
				"(The first SHA-1 certificate found has a common-name of %q.)",
			firstSHA1.Subject.CommonName,
		)
	}

	return issues
}

func checkWWW(host string) (issues Issues) {
	hasWWW := false
	if conn, err := net.DialTimeout("tcp", "www."+host+":443", dialTimeout); err == nil {
		hasWWW = true
		conn.Close()
	}

	if hasWWW {
		wwwConn, err := tls.DialWithDialer(&dialer, "tcp", "www."+host+":443", nil)
		if err != nil {
			return issues.addErrorf(
				"The www subdomain exists, but we couldn't connect to it (%q). "+
					"Since many people type this by habit, HSTS preloading would likely"+
					"cause issues for your site.",
				err,
			)
		}
		wwwConn.Close()
	}

	return issues
}

// isConnectionRefused returns true if err is an error from net/http that was
// caused because the TCP connection was refused.
func isConnectionRefused(err error) bool {
	urlErr, ok := err.(*url.Error)
	if !ok {
		return false
	}

	netErr, ok := urlErr.Err.(*net.OpError)
	if !ok {
		return false
	}

	if netErr.Op != "dial" {
		return false
	}

	syscallErr, ok := netErr.Err.(*os.SyscallError)
	if !ok {
		return false
	}

	return syscallErr.Err == syscall.ECONNREFUSED
}

func certChain(connState tls.ConnectionState) []*x509.Certificate {
	chain := connState.VerifiedChains[0]
	return chain[:len(chain)-1]
}

func findPropertyInChain(pred func(*x509.Certificate) bool, chain []*x509.Certificate) (*x509.Certificate, bool) {
	for _, cert := range chain {
		if pred(cert) {
			return cert, true
		}
	}

	return nil, false
}

func isSHA1(cert *x509.Certificate) bool {
	switch cert.SignatureAlgorithm {
	case x509.SHA1WithRSA, x509.ECDSAWithSHA1:
		return true
	default:
		return false
	}
}

func isECDSA(cert *x509.Certificate) bool {
	return cert.PublicKeyAlgorithm == x509.ECDSA
}
<commit_msg>Remove dead code from the ECDSA check.<commit_after>package hstspreload

import (
	"crypto/tls"
	"crypto/x509"
	"errors"
	"golang.org/x/net/publicsuffix"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"
)

const (
	// dialTimeout specifies the amount of time that TCP or TLS connections
	// can take to complete.
	dialTimeout = 10 * time.Second
)

// dialer is a global net.Dialer that's used whenever making TLS connections in
// order to enforce dialTimeout.
var dialer = net.Dialer{
	Timeout: dialTimeout,
}

// CheckDomain checks whether the domain passes HSTS preload
// requirements for Chromium. This includes:
//
// - Serving a single HSTS header that passes header requirements.
//
// - Using TLS settings that will not cause new problems for
// Chromium/Chrome users. (Example of a new problem: a missing intermediate certificate
// will turn an error page from overrideable to non-overridable on
// some mobile devices.)
//
// To interpret the result, see the list of conventions in the
// documentation for Issues.
func CheckDomain(domain string) (issues Issues) {
	// Check domain format issues first, since we can report something
	// useful even if the other checks fail.
	issues = combineIssues(issues, checkDomainFormat(domain))

	// We don't currently allow automatic submissions of subdomains.
	eTLD1Issues := checkEffectiveTLDPlusOne(domain)
	issues = combineIssues(issues, eTLD1Issues)

	// Start with an initial probe, and don't do the follow-up checks if
	// we can't connect.
	resp, respIssues := getResponse(domain)
	issues = combineIssues(issues, respIssues)
	if len(respIssues.Errors) == 0 {
		issues = combineIssues(issues, checkChain(certChain(*resp.TLS), domain))
		issues = combineIssues(issues, CheckResponse(*resp))

		// Skip the WWW check if the domain is not eTLD+1.
		if len(eTLD1Issues.Errors) == 0 {
			issues = combineIssues(issues, checkWWW(domain))
		}
	}

	return issues
}

func getResponse(domain string) (resp *http.Response, issues Issues) {
	redirectPrevented := errors.New("REDIRECT_PREVENTED")

	client := &http.Client{
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return redirectPrevented
		},
	}

	resp, err := client.Get("https://" + domain)
	if err != nil {
		if urlError, ok := err.(*url.Error); !ok || urlError.Err != redirectPrevented {
			return resp, issues.addErrorf(
				"Cannot connect using TLS (%q). This might be caused by an incomplete "+
					"certificate chain, which causes issues on mobile devices. "+
					"Check out your site at https://www.ssllabs.com/ssltest/",
				err,
			)
		}
	}

	return resp, issues
}

func checkDomainFormat(domain string) (issues Issues) {
	if strings.HasPrefix(domain, ".") {
		return issues.addErrorf("Domain name error: begins with `.`")
	}
	if strings.HasSuffix(domain, ".") {
		return issues.addErrorf("Domain name error: ends with `.`")
	}
	if strings.Index(domain, "..") != -1 {
		return issues.addErrorf("Domain name error: contains `..`")
	}
	if strings.Count(domain, ".") < 1 {
		return issues.addErrorf("Domain name error: must have at least two labels.")
	}

	domain = strings.ToLower(domain)
	for _, r := range domain {
		if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '-' || r == '.' {
			continue
		}

		return issues.addErrorf("Domain name error: contains invalid characters.")
	}

	return issues
}

func checkEffectiveTLDPlusOne(domain string) (issues Issues) {
	canon, err := publicsuffix.EffectiveTLDPlusOne(domain)
	if err != nil {
		return issues.addErrorf("Internal error: could not compute eTLD+1.")
	}
	if canon != domain {
		return issues.addErrorf(
			"Domain error: `%s` is not eTLD+1. Please preload `%s` instead.",
			domain,
			canon,
		)
	}

	return issues
}

// Takes the domain as argument because we may need to make more network
// connections to see if an ECDSA cert is permissible.
func checkChain(chain []*x509.Certificate, domain string) (issues Issues) {
	issues = combineIssues(issues, checkSHA1(chain))

	return issues
}

func checkSHA1(chain []*x509.Certificate) (issues Issues) {
	if firstSHA1, found := findPropertyInChain(isSHA1, chain); found {
		issues = issues.addErrorf(
			"One or more of the certificates in your certificate chain is signed with SHA-1. "+
				"This needs to be replaced. "+
				"See https://security.googleblog.com/2015/12/an-update-on-sha-1-certificates-in.html. "+
				"(The first SHA-1 certificate found has a common-name of %q.)",
			firstSHA1.Subject.CommonName,
		)
	}

	return issues
}

func checkWWW(host string) (issues Issues) {
	hasWWW := false
	if conn, err := net.DialTimeout("tcp", "www."+host+":443", dialTimeout); err == nil {
		hasWWW = true
		conn.Close()
	}

	if hasWWW {
		wwwConn, err := tls.DialWithDialer(&dialer, "tcp", "www."+host+":443", nil)
		if err != nil {
			return issues.addErrorf(
				"The www subdomain exists, but we couldn't connect to it (%q). "+
					"Since many people type this by habit, HSTS preloading would likely"+
					"cause issues for your site.",
				err,
			)
		}
		wwwConn.Close()
	}

	return issues
}

func certChain(connState tls.ConnectionState) []*x509.Certificate {
	chain := connState.VerifiedChains[0]
	return chain[:len(chain)-1]
}

func findPropertyInChain(pred func(*x509.Certificate) bool, chain []*x509.Certificate) (*x509.Certificate, bool) {
	for _, cert := range chain {
		if pred(cert) {
			return cert, true
		}
	}

	return nil, false
}

func isSHA1(cert *x509.Certificate) bool {
	switch cert.SignatureAlgorithm {
	case x509.SHA1WithRSA, x509.ECDSAWithSHA1:
		return true
	default:
		return false
	}
}
<|endoftext|>


<commit_before>package parser

import (
	"fmt"
	"os"

	"github.com/ark-lang/ark/src/util/log"

	"github.com/ark-lang/ark/src/util"
)

type IdentType int

const (
	IDENT_VARIABLE IdentType = iota
	IDENT_TYPE
	IDENT_FUNCTION
	IDENT_MODULE
)

func (v IdentType) String() string {
	switch v {
	case IDENT_FUNCTION:
		return "function"
	case IDENT_MODULE:
		return "module"
	case IDENT_TYPE:
		return "type"
	case IDENT_VARIABLE:
		return "variable"
	default:
		panic("unimplemented ident type")
	}
}

type Ident struct {
	Type  IdentType
	Value interface{}
}

type Scope struct {
	Outer  *Scope
	Idents map[string]*Ident

	UsedModules map[string]*Module
}

func newScope(outer *Scope) *Scope {
	return &Scope{
		Outer:       outer,
		Idents:      make(map[string]*Ident),
		UsedModules: make(map[string]*Module),
	}
}

var builtinScope *Scope

func init() {
	builtinScope = newScope(nil)

	for i := 0; i < len(_PrimitiveType_index); i++ {
		builtinScope.InsertType(PrimitiveType(i))
	}
}

func NewGlobalScope() *Scope {
	s := newScope(builtinScope)

	return s
}

func (v *Scope) err(err string, stuff ...interface{}) {
	// TODO: which log tag
	log.Error("parser", util.TEXT_RED+util.TEXT_BOLD+"error:"+util.TEXT_RESET+" %s\n",
		fmt.Sprintf(err, stuff...))
	os.Exit(util.EXIT_FAILURE_CODEGEN)
}

func (v *Scope) InsertIdent(value interface{}, name string, typ IdentType) *Ident {
	c := v.Idents[name]
	if c == nil {
		v.Idents[name] = &Ident{
			Type:  typ,
			Value: value,
		}
	}
	return c
}

func (v *Scope) InsertType(t Type) *Ident {
	return v.InsertIdent(t, t.TypeName(), IDENT_TYPE)
}

func (v *Scope) InsertVariable(t *Variable) *Ident {
	return v.InsertIdent(t, t.Name, IDENT_VARIABLE)
}

func (v *Scope) InsertFunction(t *Function) *Ident {
	return v.InsertIdent(t, t.Name, IDENT_FUNCTION)
}

func (v *Scope) GetIdent(name unresolvedName) *Ident {
	if len(name.moduleNames) > 0 {
		moduleName := name.moduleNames[0]

		if module, ok := v.UsedModules[moduleName]; ok {
			if r := module.GlobalScope.Idents[name.name]; r != nil {
				return r
			} else {
				v.err("could not find function `%s` in module `%s`", name.name, moduleName)
			}
		} else if v.Outer != nil {
			if module, ok := v.Outer.UsedModules[moduleName]; ok {
				if r := module.GlobalScope.Idents[name.name]; r != nil {
					return r
				} else {
					v.err("could not find function `%s` in module `%s`", name.name, moduleName)
				}
			}
		} else {
			v.err("could not find `" + moduleName + "`, are you sure it's being used in this module?\n\n    `use " + moduleName + ";`\n")
		}
	}

	if r := v.Idents[name.name]; r != nil {
		return r
	} else if v.Outer != nil {
		return v.Outer.GetIdent(name)
	}
	return nil
}
<commit_msg>err fix<commit_after>package parser

import (
	"fmt"
	"os"

	"github.com/ark-lang/ark/src/util/log"

	"github.com/ark-lang/ark/src/util"
)

type IdentType int

const (
	IDENT_VARIABLE IdentType = iota
	IDENT_TYPE
	IDENT_FUNCTION
	IDENT_MODULE
)

func (v IdentType) String() string {
	switch v {
	case IDENT_FUNCTION:
		return "function"
	case IDENT_MODULE:
		return "module"
	case IDENT_TYPE:
		return "type"
	case IDENT_VARIABLE:
		return "variable"
	default:
		panic("unimplemented ident type")
	}
}

type Ident struct {
	Type  IdentType
	Value interface{}
}

type Scope struct {
	Outer  *Scope
	Idents map[string]*Ident

	UsedModules map[string]*Module
}

func newScope(outer *Scope) *Scope {
	return &Scope{
		Outer:       outer,
		Idents:      make(map[string]*Ident),
		UsedModules: make(map[string]*Module),
	}
}

var builtinScope *Scope

func init() {
	builtinScope = newScope(nil)

	for i := 0; i < len(_PrimitiveType_index); i++ {
		builtinScope.InsertType(PrimitiveType(i))
	}
}

func NewGlobalScope() *Scope {
	s := newScope(builtinScope)

	return s
}

func (v *Scope) err(err string, stuff ...interface{}) {
	// TODO: which log tag
	log.Error("parser", util.TEXT_RED+util.TEXT_BOLD+"error:"+util.TEXT_RESET+" %s\n",
		fmt.Sprintf(err, stuff...))
	os.Exit(util.EXIT_FAILURE_CODEGEN)
}

func (v *Scope) InsertIdent(value interface{}, name string, typ IdentType) *Ident {
	c := v.Idents[name]
	if c == nil {
		v.Idents[name] = &Ident{
			Type:  typ,
			Value: value,
		}
	}
	return c
}

func (v *Scope) InsertType(t Type) *Ident {
	return v.InsertIdent(t, t.TypeName(), IDENT_TYPE)
}

func (v *Scope) InsertVariable(t *Variable) *Ident {
	return v.InsertIdent(t, t.Name, IDENT_VARIABLE)
}

func (v *Scope) InsertFunction(t *Function) *Ident {
	return v.InsertIdent(t, t.Name, IDENT_FUNCTION)
}

func (v *Scope) GetIdent(name unresolvedName) *Ident {
	if len(name.moduleNames) > 0 {
		moduleName := name.moduleNames[0]

		if module, ok := v.UsedModules[moduleName]; ok {
			if r := module.GlobalScope.Idents[name.name]; r != nil {
				return r
			} else {
				v.err("could not find identifier `%s` in module `%s`", name.name, moduleName)
			}
		} else if v.Outer != nil {
			if module, ok := v.Outer.UsedModules[moduleName]; ok {
				if r := module.GlobalScope.Idents[name.name]; r != nil {
					return r
				} else {
					v.err("could not find identifier `%s` in module `%s`", name.name, moduleName)
				}
			}
		} else {
			v.err("could not find `" + moduleName + "`, are you sure it's being used in this module?\n\n    `use " + moduleName + ";`\n")
		}
	}

	if r := v.Idents[name.name]; r != nil {
		return r
	} else if v.Outer != nil {
		return v.Outer.GetIdent(name)
	}
	return nil
}
<|endoftext|>


<commit_before>package moka

type Double interface {
	StubMethod(methodName string, args []interface{}, returnValue interface{})
	Call(methodName string, args ...interface{}) Result
}

type StrictDouble struct {
	stubs map[string]interface{}
}

func NewStrictDouble() StrictDouble {
	return StrictDouble{stubs: make(map[string]interface{})}
}

func (d StrictDouble) Call(methodName string, args ...interface{}) Result {
	return Result{d.stubs[methodName]}
}

func (d StrictDouble) StubMethod(methodName string, args []interface{}, returnValue interface{}) {
	d.stubs[methodName] = returnValue
}

type Result []interface{}

func (r Result) Get(index int) interface{} {
	return r[index]
}
<commit_msg>Result -> ReturnValues<commit_after>package moka

type Double interface {
	StubMethod(methodName string, args []interface{}, returnValue interface{})
	Call(methodName string, args ...interface{}) ReturnValues
}

type StrictDouble struct {
	stubs map[string]interface{}
}

func NewStrictDouble() StrictDouble {
	return StrictDouble{stubs: make(map[string]interface{})}
}

func (d StrictDouble) Call(methodName string, args ...interface{}) ReturnValues {
	return ReturnValues{d.stubs[methodName]}
}

func (d StrictDouble) StubMethod(methodName string, args []interface{}, returnValue interface{}) {
	d.stubs[methodName] = returnValue
}

type ReturnValues []interface{}

func (r ReturnValues) Get(index int) interface{} {
	return r[index]
}
<|endoftext|>


<commit_before>// Copyright 2010 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 html implements an HTML5-compliant tokenizer and parser.
INCOMPLETE.

Tokenization is done by creating a Tokenizer for an io.Reader r. It is the
caller's responsibility to ensure that r provides UTF-8 encoded HTML.

	z := html.NewTokenizer(r)

Given a Tokenizer z, the HTML is tokenized by repeatedly calling z.Next(),
which parses the next token and returns its type, or an error:

	for {
		tt := z.Next()
		if tt == html.ErrorToken {
			// ...
			return ...
		}
		// Process the current token.
	}

There are two APIs for retrieving the current token. The high-level API is to
call Token; the low-level API is to call Text or TagName / TagAttr. Both APIs
allow optionally calling Raw after Next but before Token, Text, TagName, or
TagAttr. In EBNF notation, the valid call sequence per token is:

	Next {Raw} [ Token | Text | TagName {TagAttr} ]

Token returns an independent data structure that completely describes a token.
Entities (such as "<") are unescaped, tag names and attribute keys are
lower-cased, and attributes are collected into a []Attribute. For example:

	for {
		if z.Next() == html.ErrorToken {
			// Returning io.EOF indicates success.
			return z.Error()
		}
		emitToken(z.Token())
	}

The low-level API performs fewer allocations and copies, but the contents of
the []byte values returned by Text, TagName and TagAttr may change on the next
call to Next. For example, to extract an HTML page's anchor text:

	depth := 0
	for {
		tt := z.Next()
		switch tt {
		case ErrorToken:
			return z.Error()
		case TextToken:
			if depth > 0 {
				// emitBytes should copy the []byte it receives,
				// if it doesn't process it immediately.
				emitBytes(z.Text())
			}
		case StartTagToken, EndTagToken:
			tn, _ := z.TagName()
			if len(tn) == 1 && tn[0] == 'a' {
				if tt == StartTag {
					depth++
				} else {
					depth--
				}
			}
		}
	}

Parsing is done by calling Parse with an io.Reader, which returns the root of
the parse tree (the document element) as a *Node. It is the caller's
responsibility to ensure that the Reader provides UTF-8 encoded HTML. For
example, to process each anchor node in depth-first order:

	doc, err := html.Parse(r)
	if err != nil {
		// ...
	}
	var f func(*html.Node)
	f = func(n *html.Node) {
		if n.Type == html.ElementNode && n.Data == "a" {
			// Do something with n...
		}
		for _, c := range n.Child {
			f(c)
		}
	}
	f(doc)

The relevant specifications include:
http://www.whatwg.org/specs/web-apps/current-work/multipage/syntax.html and
http://www.whatwg.org/specs/web-apps/current-work/multipage/tokenization.html
*/
package html

// The tokenization algorithm implemented by this package is not a line-by-line
// transliteration of the relatively verbose state-machine in the WHATWG
// specification. A more direct approach is used instead, where the program
// counter implies the state, such as whether it is tokenizing a tag or a text
// node. Specification compliance is verified by checking expected and actual
// outputs over a test suite rather than aiming for algorithmic fidelity.

// TODO(nigeltao): Does a DOM API belong in this package or a separate one?
// TODO(nigeltao): How does parsing interact with a JavaScript engine?
<commit_msg>html: fix typo in package docs.<commit_after>// Copyright 2010 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 html implements an HTML5-compliant tokenizer and parser.
INCOMPLETE.

Tokenization is done by creating a Tokenizer for an io.Reader r. It is the
caller's responsibility to ensure that r provides UTF-8 encoded HTML.

	z := html.NewTokenizer(r)

Given a Tokenizer z, the HTML is tokenized by repeatedly calling z.Next(),
which parses the next token and returns its type, or an error:

	for {
		tt := z.Next()
		if tt == html.ErrorToken {
			// ...
			return ...
		}
		// Process the current token.
	}

There are two APIs for retrieving the current token. The high-level API is to
call Token; the low-level API is to call Text or TagName / TagAttr. Both APIs
allow optionally calling Raw after Next but before Token, Text, TagName, or
TagAttr. In EBNF notation, the valid call sequence per token is:

	Next {Raw} [ Token | Text | TagName {TagAttr} ]

Token returns an independent data structure that completely describes a token.
Entities (such as "<") are unescaped, tag names and attribute keys are
lower-cased, and attributes are collected into a []Attribute. For example:

	for {
		if z.Next() == html.ErrorToken {
			// Returning io.EOF indicates success.
			return z.Error()
		}
		emitToken(z.Token())
	}

The low-level API performs fewer allocations and copies, but the contents of
the []byte values returned by Text, TagName and TagAttr may change on the next
call to Next. For example, to extract an HTML page's anchor text:

	depth := 0
	for {
		tt := z.Next()
		switch tt {
		case ErrorToken:
			return z.Error()
		case TextToken:
			if depth > 0 {
				// emitBytes should copy the []byte it receives,
				// if it doesn't process it immediately.
				emitBytes(z.Text())
			}
		case StartTagToken, EndTagToken:
			tn, _ := z.TagName()
			if len(tn) == 1 && tn[0] == 'a' {
				if tt == StartTagToken {
					depth++
				} else {
					depth--
				}
			}
		}
	}

Parsing is done by calling Parse with an io.Reader, which returns the root of
the parse tree (the document element) as a *Node. It is the caller's
responsibility to ensure that the Reader provides UTF-8 encoded HTML. For
example, to process each anchor node in depth-first order:

	doc, err := html.Parse(r)
	if err != nil {
		// ...
	}
	var f func(*html.Node)
	f = func(n *html.Node) {
		if n.Type == html.ElementNode && n.Data == "a" {
			// Do something with n...
		}
		for _, c := range n.Child {
			f(c)
		}
	}
	f(doc)

The relevant specifications include:
http://www.whatwg.org/specs/web-apps/current-work/multipage/syntax.html and
http://www.whatwg.org/specs/web-apps/current-work/multipage/tokenization.html
*/
package html

// The tokenization algorithm implemented by this package is not a line-by-line
// transliteration of the relatively verbose state-machine in the WHATWG
// specification. A more direct approach is used instead, where the program
// counter implies the state, such as whether it is tokenizing a tag or a text
// node. Specification compliance is verified by checking expected and actual
// outputs over a test suite rather than aiming for algorithmic fidelity.

// TODO(nigeltao): Does a DOM API belong in this package or a separate one?
// TODO(nigeltao): How does parsing interact with a JavaScript engine?
<|endoftext|>


<commit_before>package easyss

import (
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"net/http"
	"sync/atomic"
	"time"

	rotatelogs "github.com/lestrrat-go/file-rotatelogs"
	"github.com/nange/easypool"
	log "github.com/sirupsen/logrus"
	"github.com/txthinking/socks5"
)

const version = "1.1.2"

func init() {
	log.SetFormatter(&log.JSONFormatter{TimestampFormat: "2006-01-02 15:04:05.000"})
	log.SetLevel(log.InfoLevel)
}

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

type Statistics struct {
	BytesSend   int64
	BytesRecive int64
}

type Easyss struct {
	config        *Config
	tcpPool       easypool.Pool
	handle        *socks5.DefaultHandle
	LogFileWriter io.Writer
	stat          *Statistics

	socksServer     *socks5.Server
	httpProxyServer *http.Server
}

func New(config *Config) (*Easyss, error) {
	ss := &Easyss{
		config: config,
		handle: &socks5.DefaultHandle{},
		stat:   &Statistics{},
	}
	go ss.printStatistics()

	return ss, nil
}

func (ss *Easyss) GetLogFileFullPathName() string {
	if rl, ok := ss.LogFileWriter.(*rotatelogs.RotateLogs); ok {
		return rl.CurrentFileName()
	}
	log.Errorf("get log file name failed!")
	return ""
}

func (ss *Easyss) InitTcpPool() error {
	factory := func() (net.Conn, error) {
		return tls.Dial("tcp", fmt.Sprintf("%s:%d", ss.config.Server, ss.config.ServerPort), nil)
	}
	pconfig := &easypool.PoolConfig{
		InitialCap:  5,
		MaxCap:      30,
		MaxIdle:     5,
		Idletime:    time.Minute,
		MaxLifetime: 5 * time.Minute,
		Factory:     factory,
	}
	tcppool, err := easypool.NewHeapPool(pconfig)
	ss.tcpPool = tcppool
	return err
}

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

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

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

func (ss *Easyss) Close() {
	if ss.tcpPool != nil {
		ss.tcpPool.Close()
		ss.tcpPool = nil
	}
	if ss.httpProxyServer != nil {
		ss.httpProxyServer.Close()
		ss.httpProxyServer = nil
	}
	if ss.socksServer != nil {
		ss.socksServer.Shutdown()
		ss.socksServer = nil
	}
}

func (ss *Easyss) printStatistics() {
	for {
		select {
		case <-time.After(15 * time.Second):
			sendSize := atomic.LoadInt64(&ss.stat.BytesSend) / (1024 * 1024)
			reciveSize := atomic.LoadInt64(&ss.stat.BytesRecive) / (1024 * 1024)
			log.Infof("easyss send data size: %vMB, recive data size: %vMB", sendSize, reciveSize)
		}
	}
}
<commit_msg>*: update version<commit_after>package easyss

import (
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"net/http"
	"sync/atomic"
	"time"

	rotatelogs "github.com/lestrrat-go/file-rotatelogs"
	"github.com/nange/easypool"
	log "github.com/sirupsen/logrus"
	"github.com/txthinking/socks5"
)

const version = "1.2.0"

func init() {
	log.SetFormatter(&log.JSONFormatter{TimestampFormat: "2006-01-02 15:04:05.000"})
	log.SetLevel(log.InfoLevel)
}

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

type Statistics struct {
	BytesSend   int64
	BytesRecive int64
}

type Easyss struct {
	config        *Config
	tcpPool       easypool.Pool
	handle        *socks5.DefaultHandle
	LogFileWriter io.Writer
	stat          *Statistics

	socksServer     *socks5.Server
	httpProxyServer *http.Server
}

func New(config *Config) (*Easyss, error) {
	ss := &Easyss{
		config: config,
		handle: &socks5.DefaultHandle{},
		stat:   &Statistics{},
	}
	go ss.printStatistics()

	return ss, nil
}

func (ss *Easyss) GetLogFileFullPathName() string {
	if rl, ok := ss.LogFileWriter.(*rotatelogs.RotateLogs); ok {
		return rl.CurrentFileName()
	}
	log.Errorf("get log file name failed!")
	return ""
}

func (ss *Easyss) InitTcpPool() error {
	factory := func() (net.Conn, error) {
		return tls.Dial("tcp", fmt.Sprintf("%s:%d", ss.config.Server, ss.config.ServerPort), nil)
	}
	pconfig := &easypool.PoolConfig{
		InitialCap:  5,
		MaxCap:      30,
		MaxIdle:     5,
		Idletime:    time.Minute,
		MaxLifetime: 5 * time.Minute,
		Factory:     factory,
	}
	tcppool, err := easypool.NewHeapPool(pconfig)
	ss.tcpPool = tcppool
	return err
}

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

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

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

func (ss *Easyss) Close() {
	if ss.tcpPool != nil {
		ss.tcpPool.Close()
		ss.tcpPool = nil
	}
	if ss.httpProxyServer != nil {
		ss.httpProxyServer.Close()
		ss.httpProxyServer = nil
	}
	if ss.socksServer != nil {
		ss.socksServer.Shutdown()
		ss.socksServer = nil
	}
}

func (ss *Easyss) printStatistics() {
	for {
		select {
		case <-time.After(15 * time.Second):
			sendSize := atomic.LoadInt64(&ss.stat.BytesSend) / (1024 * 1024)
			reciveSize := atomic.LoadInt64(&ss.stat.BytesRecive) / (1024 * 1024)
			log.Infof("easyss send data size: %vMB, recive data size: %vMB", sendSize, reciveSize)
		}
	}
}
<|endoftext|>