<commit_before>//: ----------------------------------------------------------------------------
//: Copyright (C) 2017 Verizon.  All Rights Reserved.
//: All Rights Reserved
//:
//: file:    mirror.go
//: details: TODO
//: author:  Mehrdad Arshad Rad
//: date:    02/01/2017
//:
//: 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 mirror

import (
	"net"
	"syscall"
)

const (
	IPv4HLen = 20
	IPv6HLen = 40
	UDPHLen  = 8
)

type Conn struct {
	family int
	sotype int
	proto  int
	fd     int
	raddr  syscall.Sockaddr
}

func NewRawConn(raddr net.IP) (Conn, error) {
	var err error
	conn := Conn{
		sotype: syscall.SOCK_RAW,
		proto:  syscall.IPPROTO_RAW,
	}

	if ipv4 := raddr.To4(); ipv4 != nil {
		ip := [4]byte{}
		copy(ip[:], ipv4)

		conn.family = syscall.AF_INET
		conn.raddr = &syscall.SockaddrInet4{
			Port: 0,
			Addr: ip,
		}
	} else if ipv6 := raddr.To16(); ipv6 != nil {
		ip := [16]byte{}
		copy(ip[:], ipv6)

		conn.family = syscall.AF_INET6
		conn.raddr = &syscall.SockaddrInet6{
			Addr: ip,
		}

	}

	conn.fd, err = syscall.Socket(conn.family, conn.sotype, conn.proto)

	return conn, err
}

func (c *Conn) Send(b []byte) error {
	return syscall.Sendto(c.fd, b, 0, c.raddr)
}

func (c *Conn) Close(b []byte) error {
	return syscall.Close(c.fd)
}
<commit_msg>add UDP proto constant<commit_after>//: ----------------------------------------------------------------------------
//: Copyright (C) 2017 Verizon.  All Rights Reserved.
//: All Rights Reserved
//:
//: file:    mirror.go
//: details: TODO
//: author:  Mehrdad Arshad Rad
//: date:    02/01/2017
//:
//: 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 mirror

import (
	"net"
	"syscall"
)

const (
	IPv4HLen = 20
	IPv6HLen = 40
	UDPHLen  = 8
	UDPProto = 17
)

type Conn struct {
	family int
	sotype int
	proto  int
	fd     int
	raddr  syscall.Sockaddr
}

func NewRawConn(raddr net.IP) (Conn, error) {
	var err error
	conn := Conn{
		sotype: syscall.SOCK_RAW,
		proto:  syscall.IPPROTO_RAW,
	}

	if ipv4 := raddr.To4(); ipv4 != nil {
		ip := [4]byte{}
		copy(ip[:], ipv4)

		conn.family = syscall.AF_INET
		conn.raddr = &syscall.SockaddrInet4{
			Port: 0,
			Addr: ip,
		}
	} else if ipv6 := raddr.To16(); ipv6 != nil {
		ip := [16]byte{}
		copy(ip[:], ipv6)

		conn.family = syscall.AF_INET6
		conn.raddr = &syscall.SockaddrInet6{
			Addr: ip,
		}

	}

	conn.fd, err = syscall.Socket(conn.family, conn.sotype, conn.proto)

	return conn, err
}

func (c *Conn) Send(b []byte) error {
	return syscall.Sendto(c.fd, b, 0, c.raddr)
}

func (c *Conn) Close(b []byte) error {
	return syscall.Close(c.fd)
}
<|endoftext|>


<commit_before>package s3api

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"

	"github.com/golang/protobuf/jsonpb"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/iam_pb"
)

type Action string

const (
	ACTION_READ  = "Read"
	ACTION_WRITE = "Write"
	ACTION_ADMIN = "Admin"
)

type Iam interface {
	Check(f http.HandlerFunc, actions ...Action) http.HandlerFunc
}

type IdentityAccessManagement struct {
	identities []*Identity
	domain     string
}

type Identity struct {
	Name        string
	Credentials []*Credential
	Actions     []Action
}

type Credential struct {
	AccessKey string
	SecretKey string
}

func NewIdentityAccessManagement(fileName string, domain string) *IdentityAccessManagement {
	iam := &IdentityAccessManagement{
		domain: domain,
	}
	if fileName == "" {
		return iam
	}
	if err := iam.loadS3ApiConfiguration(fileName); err != nil {
		glog.Fatalf("fail to load config file %s: %v", fileName, err)
	}
	return iam
}

func (iam *IdentityAccessManagement) loadS3ApiConfiguration(fileName string) error {

	s3ApiConfiguration := &iam_pb.S3ApiConfiguration{}

	rawData, readErr := ioutil.ReadFile(fileName)
	if readErr != nil {
		glog.Warningf("fail to read %s : %v", fileName, readErr)
		return fmt.Errorf("fail to read %s : %v", fileName, readErr)
	}

	glog.V(1).Infof("maybeLoadVolumeInfo Unmarshal volume info %v", fileName)
	if err := jsonpb.Unmarshal(bytes.NewReader(rawData), s3ApiConfiguration); err != nil {
		glog.Warningf("unmarshal error: %v", err)
		return fmt.Errorf("unmarshal %s error: %v", fileName, err)
	}

	for _, ident := range s3ApiConfiguration.Identities {
		t := &Identity{
			Name:        ident.Name,
			Credentials: nil,
			Actions:     nil,
		}
		for _, action := range ident.Actions {
			t.Actions = append(t.Actions, Action(action))
		}
		for _, cred := range ident.Credentials {
			t.Credentials = append(t.Credentials, &Credential{
				AccessKey: cred.AccessKey,
				SecretKey: cred.SecretKey,
			})
		}
		iam.identities = append(iam.identities, t)
	}

	return nil
}

func (iam *IdentityAccessManagement) lookupByAccessKey(accessKey string) (identity *Identity, cred *Credential, found bool) {
	for _, ident := range iam.identities {
		for _, cred := range ident.Credentials {
			if cred.AccessKey == accessKey {
				return ident, cred, true
			}
		}
	}
	return nil, nil, false
}

func (iam *IdentityAccessManagement) Auth(f http.HandlerFunc, actions ...Action) http.HandlerFunc {

	if len(iam.identities) == 0 {
		return f
	}

	return func(w http.ResponseWriter, r *http.Request) {
		errCode := iam.authRequest(r, actions)
		if errCode == ErrNone {
			f(w, r)
			return
		}
		writeErrorResponse(w, errCode, r.URL)
	}
}

// check whether the request has valid access keys
func (iam *IdentityAccessManagement) authRequest(r *http.Request, actions []Action) ErrorCode {
	var identity *Identity
	var s3Err ErrorCode
	switch getRequestAuthType(r) {
	case authTypeStreamingSigned:
		return ErrNone
	case authTypeUnknown:
		glog.V(3).Infof("unknown auth type")
		return ErrAccessDenied
	case authTypePresignedV2, authTypeSignedV2:
		glog.V(3).Infof("v2 auth type")
		identity, s3Err = iam.isReqAuthenticatedV2(r)
	case authTypeSigned, authTypePresigned:
		glog.V(3).Infof("v4 auth type")
		identity, s3Err = iam.reqSignatureV4Verify(r)
	case authTypePostPolicy:
		return ErrNotImplemented;
	case authTypeJWT:
		return ErrNotImplemented;
	case authTypeAnonymous:
		return ErrNotImplemented
	}

	glog.V(3).Infof("auth error: %v", s3Err)
	if s3Err != ErrNone {
		return s3Err
	}

	glog.V(3).Infof("user name: %v actions: %v", identity.Name, identity.Actions)

	if !identity.canDo(actions) {
		return ErrAccessDenied
	}

	return ErrNone

}

func (identity *Identity) canDo(actions []Action) bool {
	for _, a := range identity.Actions {
		for _, b := range actions {
			if a == b {
				return true
			}
		}
	}
	return false
}
<commit_msg>s3: deny anonymous type<commit_after>package s3api

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net/http"

	"github.com/golang/protobuf/jsonpb"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/iam_pb"
)

type Action string

const (
	ACTION_READ  = "Read"
	ACTION_WRITE = "Write"
	ACTION_ADMIN = "Admin"
)

type Iam interface {
	Check(f http.HandlerFunc, actions ...Action) http.HandlerFunc
}

type IdentityAccessManagement struct {
	identities []*Identity
	domain     string
}

type Identity struct {
	Name        string
	Credentials []*Credential
	Actions     []Action
}

type Credential struct {
	AccessKey string
	SecretKey string
}

func NewIdentityAccessManagement(fileName string, domain string) *IdentityAccessManagement {
	iam := &IdentityAccessManagement{
		domain: domain,
	}
	if fileName == "" {
		return iam
	}
	if err := iam.loadS3ApiConfiguration(fileName); err != nil {
		glog.Fatalf("fail to load config file %s: %v", fileName, err)
	}
	return iam
}

func (iam *IdentityAccessManagement) loadS3ApiConfiguration(fileName string) error {

	s3ApiConfiguration := &iam_pb.S3ApiConfiguration{}

	rawData, readErr := ioutil.ReadFile(fileName)
	if readErr != nil {
		glog.Warningf("fail to read %s : %v", fileName, readErr)
		return fmt.Errorf("fail to read %s : %v", fileName, readErr)
	}

	glog.V(1).Infof("maybeLoadVolumeInfo Unmarshal volume info %v", fileName)
	if err := jsonpb.Unmarshal(bytes.NewReader(rawData), s3ApiConfiguration); err != nil {
		glog.Warningf("unmarshal error: %v", err)
		return fmt.Errorf("unmarshal %s error: %v", fileName, err)
	}

	for _, ident := range s3ApiConfiguration.Identities {
		t := &Identity{
			Name:        ident.Name,
			Credentials: nil,
			Actions:     nil,
		}
		for _, action := range ident.Actions {
			t.Actions = append(t.Actions, Action(action))
		}
		for _, cred := range ident.Credentials {
			t.Credentials = append(t.Credentials, &Credential{
				AccessKey: cred.AccessKey,
				SecretKey: cred.SecretKey,
			})
		}
		iam.identities = append(iam.identities, t)
	}

	return nil
}

func (iam *IdentityAccessManagement) lookupByAccessKey(accessKey string) (identity *Identity, cred *Credential, found bool) {
	for _, ident := range iam.identities {
		for _, cred := range ident.Credentials {
			if cred.AccessKey == accessKey {
				return ident, cred, true
			}
		}
	}
	return nil, nil, false
}

func (iam *IdentityAccessManagement) Auth(f http.HandlerFunc, actions ...Action) http.HandlerFunc {

	if len(iam.identities) == 0 {
		return f
	}

	return func(w http.ResponseWriter, r *http.Request) {
		errCode := iam.authRequest(r, actions)
		if errCode == ErrNone {
			f(w, r)
			return
		}
		writeErrorResponse(w, errCode, r.URL)
	}
}

// check whether the request has valid access keys
func (iam *IdentityAccessManagement) authRequest(r *http.Request, actions []Action) ErrorCode {
	var identity *Identity
	var s3Err ErrorCode
	switch getRequestAuthType(r) {
	case authTypeStreamingSigned:
		return ErrNone
	case authTypeUnknown:
		glog.V(3).Infof("unknown auth type")
		return ErrAccessDenied
	case authTypePresignedV2, authTypeSignedV2:
		glog.V(3).Infof("v2 auth type")
		identity, s3Err = iam.isReqAuthenticatedV2(r)
	case authTypeSigned, authTypePresigned:
		glog.V(3).Infof("v4 auth type")
		identity, s3Err = iam.reqSignatureV4Verify(r)
	case authTypePostPolicy:
		glog.V(3).Infof("post policy auth type")
		return ErrNotImplemented
	case authTypeJWT:
		glog.V(3).Infof("jwt auth type")
		return ErrNotImplemented
	case authTypeAnonymous:
		return ErrAccessDenied
	default:
		return ErrNotImplemented
	}

	glog.V(3).Infof("auth error: %v", s3Err)
	if s3Err != ErrNone {
		return s3Err
	}

	glog.V(3).Infof("user name: %v actions: %v", identity.Name, identity.Actions)

	if !identity.canDo(actions) {
		return ErrAccessDenied
	}

	return ErrNone

}

func (identity *Identity) canDo(actions []Action) bool {
	for _, a := range identity.Actions {
		for _, b := range actions {
			if a == b {
				return true
			}
		}
	}
	return false
}
<|endoftext|>


<commit_before>package topology

import (
	"errors"
	"fmt"
	"math/rand"
	"sync"
	"time"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/storage"
)

// mapping from volume to its locations, inverted from server to volume
type VolumeLayout struct {
	rp               *storage.ReplicaPlacement
	ttl              *storage.TTL
	vid2location     map[storage.VolumeId]*VolumeLocationList
	writables        []storage.VolumeId        // transient array of writable volume id
	readonlyVolumes  map[storage.VolumeId]bool // transient set of readonly volumes
	oversizedVolumes map[storage.VolumeId]bool // set of oversized volumes
	volumeSizeLimit  uint64
	accessLock       sync.RWMutex
}

type VolumeLayoutStats struct {
	TotalSize uint64
	UsedSize  uint64
	FileCount uint64
}

func NewVolumeLayout(rp *storage.ReplicaPlacement, ttl *storage.TTL, volumeSizeLimit uint64) *VolumeLayout {
	return &VolumeLayout{
		rp:               rp,
		ttl:              ttl,
		vid2location:     make(map[storage.VolumeId]*VolumeLocationList),
		writables:        *new([]storage.VolumeId),
		readonlyVolumes:  make(map[storage.VolumeId]bool),
		oversizedVolumes: make(map[storage.VolumeId]bool),
		volumeSizeLimit:  volumeSizeLimit,
	}
}

func (vl *VolumeLayout) String() string {
	return fmt.Sprintf("rp:%v, ttl:%v, vid2location:%v, writables:%v, volumeSizeLimit:%v", vl.rp, vl.ttl, vl.vid2location, vl.writables, vl.volumeSizeLimit)
}

func (vl *VolumeLayout) RegisterVolume(v *storage.VolumeInfo, dn *DataNode) {
	vl.accessLock.Lock()
	defer vl.accessLock.Unlock()

	if _, ok := vl.vid2location[v.Id]; !ok {
		vl.vid2location[v.Id] = NewVolumeLocationList()
	}
	vl.vid2location[v.Id].Set(dn)
	glog.V(4).Infof("volume %d added to %s len %d copy %d", v.Id, dn.Id(), vl.vid2location[v.Id].Length(), v.ReplicaPlacement.GetCopyCount())
	for _, dn := range vl.vid2location[v.Id].list {
		if v_info, err := dn.GetVolumesById(v.Id); err == nil {
			if v_info.ReadOnly {
				glog.V(3).Infof("vid %d removed from writable", v.Id)
				vl.removeFromWritable(v.Id)
				vl.readonlyVolumes[v.Id] = true
				return
			} else {
				delete(vl.readonlyVolumes, v.Id)
			}
		} else {
			glog.V(3).Infof("vid %d removed from writable", v.Id)
			vl.removeFromWritable(v.Id)
			delete(vl.readonlyVolumes, v.Id)
			return
		}
	}
	if vl.vid2location[v.Id].Length() == vl.rp.GetCopyCount() && vl.isWritable(v) {
		if _, ok := vl.oversizedVolumes[v.Id]; !ok {
			vl.addToWritable(v.Id)
		}
	} else {
		vl.rememberOversizedVolumne(v)
		vl.removeFromWritable(v.Id)
	}
}

func (vl *VolumeLayout) rememberOversizedVolumne(v *storage.VolumeInfo) {
	if vl.isOversized(v) {
		vl.oversizedVolumes[v.Id] = true
	}
}

func (vl *VolumeLayout) UnRegisterVolume(v *storage.VolumeInfo, dn *DataNode) {
	vl.accessLock.Lock()
	defer vl.accessLock.Unlock()

	vl.removeFromWritable(v.Id)
	delete(vl.vid2location, v.Id)
}

func (vl *VolumeLayout) addToWritable(vid storage.VolumeId) {
	for _, id := range vl.writables {
		if vid == id {
			return
		}
	}
	vl.writables = append(vl.writables, vid)
}

func (vl *VolumeLayout) isOversized(v *storage.VolumeInfo) bool {
	return uint64(v.Size) >= vl.volumeSizeLimit
}

func (vl *VolumeLayout) isWritable(v *storage.VolumeInfo) bool {
	return !vl.isOversized(v) &&
		v.Version == storage.CurrentVersion &&
		!v.ReadOnly
}

func (vl *VolumeLayout) isEmpty() bool {
	vl.accessLock.RLock()
	defer vl.accessLock.RUnlock()

	return len(vl.vid2location) == 0
}

func (vl *VolumeLayout) Lookup(vid storage.VolumeId) []*DataNode {
	vl.accessLock.RLock()
	defer vl.accessLock.RUnlock()

	if location := vl.vid2location[vid]; location != nil {
		return location.list
	}
	return nil
}

func (vl *VolumeLayout) ListVolumeServers() (nodes []*DataNode) {
	vl.accessLock.RLock()
	defer vl.accessLock.RUnlock()

	for _, location := range vl.vid2location {
		nodes = append(nodes, location.list...)
	}
	return
}

func (vl *VolumeLayout) PickForWrite(count uint64, option *VolumeGrowOption) (*storage.VolumeId, uint64, *VolumeLocationList, error) {
	vl.accessLock.RLock()
	defer vl.accessLock.RUnlock()

	len_writers := len(vl.writables)
	if len_writers <= 0 {
		glog.V(0).Infoln("No more writable volumes!")
		return nil, 0, nil, errors.New("No more writable volumes!")
	}
	if option.DataCenter == "" {
		vid := vl.writables[rand.Intn(len_writers)]
		locationList := vl.vid2location[vid]
		if locationList != nil {
			return &vid, count, locationList, nil
		}
		return nil, 0, nil, errors.New("Strangely vid " + vid.String() + " is on no machine!")
	}
	var vid storage.VolumeId
	var locationList *VolumeLocationList
	counter := 0
	for _, v := range vl.writables {
		volumeLocationList := vl.vid2location[v]
		for _, dn := range volumeLocationList.list {
			if dn.GetDataCenter().Id() == NodeId(option.DataCenter) {
				if option.Rack != "" && dn.GetRack().Id() != NodeId(option.Rack) {
					continue
				}
				if option.DataNode != "" && dn.Id() != NodeId(option.DataNode) {
					continue
				}
				counter++
				if rand.Intn(counter) < 1 {
					vid, locationList = v, volumeLocationList
				}
			}
		}
	}
	return &vid, count, locationList, nil
}

func (vl *VolumeLayout) GetActiveVolumeCount(option *VolumeGrowOption) int {
	vl.accessLock.RLock()
	defer vl.accessLock.RUnlock()

	if option.DataCenter == "" {
		return len(vl.writables)
	}
	counter := 0
	for _, v := range vl.writables {
		for _, dn := range vl.vid2location[v].list {
			if dn.GetDataCenter().Id() == NodeId(option.DataCenter) {
				if option.Rack != "" && dn.GetRack().Id() != NodeId(option.Rack) {
					continue
				}
				if option.DataNode != "" && dn.Id() != NodeId(option.DataNode) {
					continue
				}
				counter++
			}
		}
	}
	return counter
}

func (vl *VolumeLayout) removeFromWritable(vid storage.VolumeId) bool {
	toDeleteIndex := -1
	for k, id := range vl.writables {
		if id == vid {
			toDeleteIndex = k
			break
		}
	}
	if toDeleteIndex >= 0 {
		glog.V(0).Infoln("Volume", vid, "becomes unwritable")
		vl.writables = append(vl.writables[0:toDeleteIndex], vl.writables[toDeleteIndex+1:]...)
		return true
	}
	return false
}
func (vl *VolumeLayout) setVolumeWritable(vid storage.VolumeId) bool {
	for _, v := range vl.writables {
		if v == vid {
			return false
		}
	}
	glog.V(0).Infoln("Volume", vid, "becomes writable")
	vl.writables = append(vl.writables, vid)
	return true
}

func (vl *VolumeLayout) SetVolumeUnavailable(dn *DataNode, vid storage.VolumeId) bool {
	vl.accessLock.Lock()
	defer vl.accessLock.Unlock()

	if location, ok := vl.vid2location[vid]; ok {
		if location.Remove(dn) {
			if location.Length() < vl.rp.GetCopyCount() {
				glog.V(0).Infoln("Volume", vid, "has", location.Length(), "replica, less than required", vl.rp.GetCopyCount())
				return vl.removeFromWritable(vid)
			}
		}
	}
	return false
}
func (vl *VolumeLayout) SetVolumeAvailable(dn *DataNode, vid storage.VolumeId) bool {
	vl.accessLock.Lock()
	defer vl.accessLock.Unlock()

	vl.vid2location[vid].Set(dn)
	if vl.vid2location[vid].Length() >= vl.rp.GetCopyCount() {
		return vl.setVolumeWritable(vid)
	}
	return false
}

func (vl *VolumeLayout) SetVolumeCapacityFull(vid storage.VolumeId) bool {
	vl.accessLock.Lock()
	defer vl.accessLock.Unlock()

	// glog.V(0).Infoln("Volume", vid, "reaches full capacity.")
	return vl.removeFromWritable(vid)
}

func (vl *VolumeLayout) ToMap() map[string]interface{} {
	m := make(map[string]interface{})
	m["replication"] = vl.rp.String()
	m["ttl"] = vl.ttl.String()
	m["writables"] = vl.writables
	//m["locations"] = vl.vid2location
	return m
}

func (vl *VolumeLayout) Stats() *VolumeLayoutStats {
	vl.accessLock.RLock()
	defer vl.accessLock.RUnlock()

	ret := &VolumeLayoutStats{}

	freshThreshold := time.Now().Unix() - 60

	for vid, vll := range vl.vid2location {
		size, fileCount := vll.Stats(vid, freshThreshold)
		ret.FileCount += uint64(fileCount)
		ret.UsedSize += size
		if vl.readonlyVolumes[vid] {
			ret.TotalSize += size
		} else {
			ret.TotalSize += vl.volumeSizeLimit
		}
	}

	return ret
}
<commit_msg>adjust log<commit_after>package topology

import (
	"errors"
	"fmt"
	"math/rand"
	"sync"
	"time"

	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/storage"
)

// mapping from volume to its locations, inverted from server to volume
type VolumeLayout struct {
	rp               *storage.ReplicaPlacement
	ttl              *storage.TTL
	vid2location     map[storage.VolumeId]*VolumeLocationList
	writables        []storage.VolumeId        // transient array of writable volume id
	readonlyVolumes  map[storage.VolumeId]bool // transient set of readonly volumes
	oversizedVolumes map[storage.VolumeId]bool // set of oversized volumes
	volumeSizeLimit  uint64
	accessLock       sync.RWMutex
}

type VolumeLayoutStats struct {
	TotalSize uint64
	UsedSize  uint64
	FileCount uint64
}

func NewVolumeLayout(rp *storage.ReplicaPlacement, ttl *storage.TTL, volumeSizeLimit uint64) *VolumeLayout {
	return &VolumeLayout{
		rp:               rp,
		ttl:              ttl,
		vid2location:     make(map[storage.VolumeId]*VolumeLocationList),
		writables:        *new([]storage.VolumeId),
		readonlyVolumes:  make(map[storage.VolumeId]bool),
		oversizedVolumes: make(map[storage.VolumeId]bool),
		volumeSizeLimit:  volumeSizeLimit,
	}
}

func (vl *VolumeLayout) String() string {
	return fmt.Sprintf("rp:%v, ttl:%v, vid2location:%v, writables:%v, volumeSizeLimit:%v", vl.rp, vl.ttl, vl.vid2location, vl.writables, vl.volumeSizeLimit)
}

func (vl *VolumeLayout) RegisterVolume(v *storage.VolumeInfo, dn *DataNode) {
	vl.accessLock.Lock()
	defer vl.accessLock.Unlock()

	if _, ok := vl.vid2location[v.Id]; !ok {
		vl.vid2location[v.Id] = NewVolumeLocationList()
	}
	vl.vid2location[v.Id].Set(dn)
	// glog.V(4).Infof("volume %d added to %s len %d copy %d", v.Id, dn.Id(), vl.vid2location[v.Id].Length(), v.ReplicaPlacement.GetCopyCount())
	for _, dn := range vl.vid2location[v.Id].list {
		if v_info, err := dn.GetVolumesById(v.Id); err == nil {
			if v_info.ReadOnly {
				glog.V(3).Infof("vid %d removed from writable", v.Id)
				vl.removeFromWritable(v.Id)
				vl.readonlyVolumes[v.Id] = true
				return
			} else {
				delete(vl.readonlyVolumes, v.Id)
			}
		} else {
			glog.V(3).Infof("vid %d removed from writable", v.Id)
			vl.removeFromWritable(v.Id)
			delete(vl.readonlyVolumes, v.Id)
			return
		}
	}
	if vl.vid2location[v.Id].Length() == vl.rp.GetCopyCount() && vl.isWritable(v) {
		if _, ok := vl.oversizedVolumes[v.Id]; !ok {
			vl.addToWritable(v.Id)
		}
	} else {
		vl.rememberOversizedVolumne(v)
		vl.removeFromWritable(v.Id)
	}
}

func (vl *VolumeLayout) rememberOversizedVolumne(v *storage.VolumeInfo) {
	if vl.isOversized(v) {
		vl.oversizedVolumes[v.Id] = true
	}
}

func (vl *VolumeLayout) UnRegisterVolume(v *storage.VolumeInfo, dn *DataNode) {
	vl.accessLock.Lock()
	defer vl.accessLock.Unlock()

	vl.removeFromWritable(v.Id)
	delete(vl.vid2location, v.Id)
}

func (vl *VolumeLayout) addToWritable(vid storage.VolumeId) {
	for _, id := range vl.writables {
		if vid == id {
			return
		}
	}
	vl.writables = append(vl.writables, vid)
}

func (vl *VolumeLayout) isOversized(v *storage.VolumeInfo) bool {
	return uint64(v.Size) >= vl.volumeSizeLimit
}

func (vl *VolumeLayout) isWritable(v *storage.VolumeInfo) bool {
	return !vl.isOversized(v) &&
		v.Version == storage.CurrentVersion &&
		!v.ReadOnly
}

func (vl *VolumeLayout) isEmpty() bool {
	vl.accessLock.RLock()
	defer vl.accessLock.RUnlock()

	return len(vl.vid2location) == 0
}

func (vl *VolumeLayout) Lookup(vid storage.VolumeId) []*DataNode {
	vl.accessLock.RLock()
	defer vl.accessLock.RUnlock()

	if location := vl.vid2location[vid]; location != nil {
		return location.list
	}
	return nil
}

func (vl *VolumeLayout) ListVolumeServers() (nodes []*DataNode) {
	vl.accessLock.RLock()
	defer vl.accessLock.RUnlock()

	for _, location := range vl.vid2location {
		nodes = append(nodes, location.list...)
	}
	return
}

func (vl *VolumeLayout) PickForWrite(count uint64, option *VolumeGrowOption) (*storage.VolumeId, uint64, *VolumeLocationList, error) {
	vl.accessLock.RLock()
	defer vl.accessLock.RUnlock()

	len_writers := len(vl.writables)
	if len_writers <= 0 {
		glog.V(0).Infoln("No more writable volumes!")
		return nil, 0, nil, errors.New("No more writable volumes!")
	}
	if option.DataCenter == "" {
		vid := vl.writables[rand.Intn(len_writers)]
		locationList := vl.vid2location[vid]
		if locationList != nil {
			return &vid, count, locationList, nil
		}
		return nil, 0, nil, errors.New("Strangely vid " + vid.String() + " is on no machine!")
	}
	var vid storage.VolumeId
	var locationList *VolumeLocationList
	counter := 0
	for _, v := range vl.writables {
		volumeLocationList := vl.vid2location[v]
		for _, dn := range volumeLocationList.list {
			if dn.GetDataCenter().Id() == NodeId(option.DataCenter) {
				if option.Rack != "" && dn.GetRack().Id() != NodeId(option.Rack) {
					continue
				}
				if option.DataNode != "" && dn.Id() != NodeId(option.DataNode) {
					continue
				}
				counter++
				if rand.Intn(counter) < 1 {
					vid, locationList = v, volumeLocationList
				}
			}
		}
	}
	return &vid, count, locationList, nil
}

func (vl *VolumeLayout) GetActiveVolumeCount(option *VolumeGrowOption) int {
	vl.accessLock.RLock()
	defer vl.accessLock.RUnlock()

	if option.DataCenter == "" {
		return len(vl.writables)
	}
	counter := 0
	for _, v := range vl.writables {
		for _, dn := range vl.vid2location[v].list {
			if dn.GetDataCenter().Id() == NodeId(option.DataCenter) {
				if option.Rack != "" && dn.GetRack().Id() != NodeId(option.Rack) {
					continue
				}
				if option.DataNode != "" && dn.Id() != NodeId(option.DataNode) {
					continue
				}
				counter++
			}
		}
	}
	return counter
}

func (vl *VolumeLayout) removeFromWritable(vid storage.VolumeId) bool {
	toDeleteIndex := -1
	for k, id := range vl.writables {
		if id == vid {
			toDeleteIndex = k
			break
		}
	}
	if toDeleteIndex >= 0 {
		glog.V(0).Infoln("Volume", vid, "becomes unwritable")
		vl.writables = append(vl.writables[0:toDeleteIndex], vl.writables[toDeleteIndex+1:]...)
		return true
	}
	return false
}
func (vl *VolumeLayout) setVolumeWritable(vid storage.VolumeId) bool {
	for _, v := range vl.writables {
		if v == vid {
			return false
		}
	}
	glog.V(0).Infoln("Volume", vid, "becomes writable")
	vl.writables = append(vl.writables, vid)
	return true
}

func (vl *VolumeLayout) SetVolumeUnavailable(dn *DataNode, vid storage.VolumeId) bool {
	vl.accessLock.Lock()
	defer vl.accessLock.Unlock()

	if location, ok := vl.vid2location[vid]; ok {
		if location.Remove(dn) {
			if location.Length() < vl.rp.GetCopyCount() {
				glog.V(0).Infoln("Volume", vid, "has", location.Length(), "replica, less than required", vl.rp.GetCopyCount())
				return vl.removeFromWritable(vid)
			}
		}
	}
	return false
}
func (vl *VolumeLayout) SetVolumeAvailable(dn *DataNode, vid storage.VolumeId) bool {
	vl.accessLock.Lock()
	defer vl.accessLock.Unlock()

	vl.vid2location[vid].Set(dn)
	if vl.vid2location[vid].Length() >= vl.rp.GetCopyCount() {
		return vl.setVolumeWritable(vid)
	}
	return false
}

func (vl *VolumeLayout) SetVolumeCapacityFull(vid storage.VolumeId) bool {
	vl.accessLock.Lock()
	defer vl.accessLock.Unlock()

	// glog.V(0).Infoln("Volume", vid, "reaches full capacity.")
	return vl.removeFromWritable(vid)
}

func (vl *VolumeLayout) ToMap() map[string]interface{} {
	m := make(map[string]interface{})
	m["replication"] = vl.rp.String()
	m["ttl"] = vl.ttl.String()
	m["writables"] = vl.writables
	//m["locations"] = vl.vid2location
	return m
}

func (vl *VolumeLayout) Stats() *VolumeLayoutStats {
	vl.accessLock.RLock()
	defer vl.accessLock.RUnlock()

	ret := &VolumeLayoutStats{}

	freshThreshold := time.Now().Unix() - 60

	for vid, vll := range vl.vid2location {
		size, fileCount := vll.Stats(vid, freshThreshold)
		ret.FileCount += uint64(fileCount)
		ret.UsedSize += size
		if vl.readonlyVolumes[vid] {
			ret.TotalSize += size
		} else {
			ret.TotalSize += vl.volumeSizeLimit
		}
	}

	return ret
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"github.com/thoj/go-ircevent"
	"io/ioutil"
	"regexp"
	"time"
	"strconv"
	"os"
)

type Config struct {
	Nick     string
	User     string
	Server   string
	Channels []string
	Admins   []string
}

type UserData struct {
	Date time.Time
}

var (
	users     map[string]UserData
	admins    []string
	cmdRe     *regexp.Regexp
	chatCmdRe *regexp.Regexp
	dateRe    *regexp.Regexp
)

func loadConfig(path string) Config {
	var c Config
	b, err := ioutil.ReadFile(path)
	if err != nil {
		panic(err)
	}
	err = json.Unmarshal(b, &c)
	if err != nil {
		panic(err)
	}
	return c
}

func isAdmin(nick string) bool {
	for _, s := range admins {
		if nick == s {
			return true
		}
	}
	return false
}

func setDate(user string, date time.Time) {
	users[user] = UserData{date}
	save()
}

func save() {
	data, _ := json.Marshal(users)
	ioutil.WriteFile("users.json", data, os.FileMode(0644))
}

func load() {
	b, err := ioutil.ReadFile("users.json")
	if err != nil {
		fmt.Println("Error reading users file")
		return
	}
	err = json.Unmarshal(b, &users)
	if err != nil {
		fmt.Println("Error parsing users file")
	}
}

var funcs = map[string]func(string, string) string{
	"hello": func(nick, data string) string {
		return "Hello, " + nick
	},
	"set": func(nick, data string) string {
		p := dateRe.FindStringSubmatch(data)
		if p == nil {
			return "Data format is yyyy-mm-dd"
		}
		dateStr := p[1]
		user := p[2]
		date, err := time.Parse("2006-01-02", dateStr)
		if err != nil {
			return err.Error()
		}
		if user == "" {
			user = nick
		}
		if user == nick || isAdmin(nick) {
			setDate(user, date)
			return "Counter of user " + user + " updated"
		} else {
			return "Sorry, you can't change other people's counters"
		}
	},
	"reset": func(nick, data string) string {
		user := data
		if user == "" {
			user = nick
		}
		now := time.Now()
		if user == nick || isAdmin(nick) {
			setDate(user, time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC))
			return "Counter for user " + user + " updated"
		}
		return "Sorry, you can't change other people's counters"
	},
	"get": func(nick, data string) string {
		user := data
		if user == "" {
			user = nick
		}
		if u, ok := users[user]; ok {
			dur := time.Now().Sub(u.Date)
			days := int(dur.Hours() / 24)
			var dstr string
			if days == 1 {
				dstr = " day"
			} else {
				dstr = " days"
			}
			return user + " is winning for " + strconv.Itoa(days) + dstr
		}
		return "Counter not found for user " + user
	},
	"help": func(nick, data string) string {
		return "Commands: help, set, get, reset, hello"
	},
}

func onMessage(bot *irc.Connection, from, to, message string) {
	fmt.Printf("%s => %s: %s\n", from, to, message)
	re := cmdRe
	answer := ""
	answerTo := from
	if to[0] == '#' {
		answer = from + ": "
		re = chatCmdRe
		answerTo = to
	}
	m := re.FindStringSubmatch(message)
	if m != nil {
		cmd, data := m[1], m[2]
		fmt.Println("Cmd:", cmd, "Data:", data)
		if funcs[cmd] != nil {
			answer += funcs[cmd](from, data)
		} else {
			answer += "Command not found"
		}
		bot.Privmsg(answerTo, answer)
	}
}

func compileRegex(nick string) {
	cmdReStr := "(\\w+)(?:\\s(.+))?" //regex for command in private
	chatCmdReStr := "^" + nick + "[:,]\\s" + cmdReStr
	cmdRe, _ = regexp.Compile(cmdReStr)
	chatCmdRe, _ = regexp.Compile(chatCmdReStr)
	dateRe, _ = regexp.Compile(`(\d\d\d\d-\d?\d-\d?\d)(?:\s(\w+))?`)
}

func main() {
	config := loadConfig("conf.json")
	admins = config.Admins
	users = make(map[string]UserData)
	load()
	compileRegex(config.Nick)
	bot := irc.IRC(config.Nick, config.User)
	bot.Connect(config.Server)
	for _, ch := range config.Channels {
		bot.Join(ch)
	}
	bot.AddCallback("PRIVMSG", func(e *irc.Event) {
		onMessage(bot, e.Nick, e.Arguments[0], e.Message)
	})
	bot.Loop()
}
<commit_msg>Commands are now case-insensitive<commit_after>// Copyright 2013 Denis Pobedrya <denis.pobedrya@gmail.com> All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package main

import (
	"encoding/json"
	"fmt"
	"github.com/thoj/go-ircevent"
	"io/ioutil"
	"regexp"
	"time"
	"strconv"
	"os"
	"strings"
)

type Config struct {
	Nick     string
	User     string
	Server   string
	Channels []string
	Admins   []string
}

type UserData struct {
	Date time.Time
}

var (
	users     map[string]UserData
	admins    []string
	cmdRe     *regexp.Regexp
	chatCmdRe *regexp.Regexp
	dateRe    *regexp.Regexp
)

func loadConfig(path string) Config {
	var c Config
	b, err := ioutil.ReadFile(path)
	if err != nil {
		panic(err)
	}
	err = json.Unmarshal(b, &c)
	if err != nil {
		panic(err)
	}
	return c
}

func isAdmin(nick string) bool {
	for _, s := range admins {
		if nick == s {
			return true
		}
	}
	return false
}

func setDate(user string, date time.Time) {
	users[user] = UserData{date}
	save()
}

func save() {
	data, _ := json.Marshal(users)
	ioutil.WriteFile("users.json", data, os.FileMode(0644))
}

func load() {
	b, err := ioutil.ReadFile("users.json")
	if err != nil {
		fmt.Println("Error reading users file")
		return
	}
	err = json.Unmarshal(b, &users)
	if err != nil {
		fmt.Println("Error parsing users file")
	}
}

var funcs = map[string]func(string, string) string{
	"hello": func(nick, data string) string {
		return "Hello, " + nick
	},
	"set": func(nick, data string) string {
		p := dateRe.FindStringSubmatch(data)
		if p == nil {
			return "Data format is yyyy-mm-dd"
		}
		dateStr := p[1]
		user := p[2]
		date, err := time.Parse("2006-01-02", dateStr)
		if err != nil {
			return err.Error()
		}
		if user == "" {
			user = nick
		}
		if user == nick || isAdmin(nick) {
			setDate(user, date)
			return "Counter of user " + user + " updated"
		} else {
			return "Sorry, you can't change other people's counters"
		}
	},
	"reset": func(nick, data string) string {
		user := data
		if user == "" {
			user = nick
		}
		now := time.Now()
		if user == nick || isAdmin(nick) {
			setDate(user, time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC))
			return "Counter for user " + user + " updated"
		}
		return "Sorry, you can't change other people's counters"
	},
	"get": func(nick, data string) string {
		user := data
		if user == "" {
			user = nick
		}
		if u, ok := users[user]; ok {
			dur := time.Now().Sub(u.Date)
			days := int(dur.Hours() / 24)
			var dstr string
			if days == 1 {
				dstr = " day"
			} else {
				dstr = " days"
			}
			return user + " is winning for " + strconv.Itoa(days) + dstr
		}
		return "Counter not found for user " + user
	},
	"help": func(nick, data string) string {
		return "Commands: help, set, get, reset, hello"
	},
}

func onMessage(bot *irc.Connection, from, to, message string) {
	fmt.Printf("%s => %s: %s\n", from, to, message)
	re := cmdRe
	answer := ""
	answerTo := from
	if to[0] == '#' {
		answer = from + ": "
		re = chatCmdRe
		answerTo = to
	}
	m := re.FindStringSubmatch(message)
	if m != nil {
		cmd, data := strings.ToLower(m[1]), m[2]
		fmt.Println("Cmd:", cmd, "Data:", data)
		if funcs[cmd] != nil {
			answer += funcs[cmd](from, data)
		} else {
			answer += "Command not found"
		}
		bot.Privmsg(answerTo, answer)
	}
}

func compileRegex(nick string) {
	cmdReStr := "(\\w+)(?:\\s(.+))?" //regex for command in private
	chatCmdReStr := "^" + nick + "[:,]\\s" + cmdReStr
	cmdRe, _ = regexp.Compile(cmdReStr)
	chatCmdRe, _ = regexp.Compile(chatCmdReStr)
	dateRe, _ = regexp.Compile(`(\d\d\d\d-\d?\d-\d?\d)(?:\s(\w+))?`)
}

func main() {
	config := loadConfig("conf.json")
	admins = config.Admins
	users = make(map[string]UserData)
	load()
	compileRegex(config.Nick)
	bot := irc.IRC(config.Nick, config.User)
	bot.Connect(config.Server)
	for _, ch := range config.Channels {
		bot.Join(ch)
	}
	bot.AddCallback("PRIVMSG", func(e *irc.Event) {
		onMessage(bot, e.Nick, e.Arguments[0], e.Message)
	})
	bot.Loop()
}
<|endoftext|>


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

package blackbox

import (
	"context"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"syscall"
	"testing"
	"time"

	config "github.com/coreos/ignition/config/v3_0_experimental"
	"github.com/coreos/ignition/tests/register"
	"github.com/coreos/ignition/tests/types"

	// Register the tests
	_ "github.com/coreos/ignition/tests/registry"

	// UUID generation tool
	"github.com/pborman/uuid"
)

var (
	// testTimeout controls how long a given test is allowed to run before being
	// cancelled.
	testTimeout = time.Second * 60
	// somewhat of an abuse of contexts but go's got our hands tied
	killContext = context.TODO()

	// flag for listing all subtests that would be run without running them
	listSubtests = false
)

func TestMain(m *testing.M) {
	interruptChan := make(chan os.Signal, 3)
	signal.Notify(interruptChan, os.Interrupt, syscall.SIGTERM)
	tmp, killCancel := context.WithCancel(context.Background())
	killContext = tmp
	go func() {
		for {
			sig := <-interruptChan
			switch sig {
			case os.Interrupt, syscall.SIGTERM:
				killCancel()
			}
		}
	}()

	flag.BoolVar(&listSubtests, "list", false, "list tests that would be run without running them")
	flag.Parse()

	if !listSubtests {
		httpServer := &HTTPServer{}
		httpServer.Start()
		tftpServer := &TFTPServer{}
		tftpServer.Start()
	}
	os.Exit(m.Run())
}

func TestIgnitionBlackBox(t *testing.T) {
	for _, test := range register.Tests[register.PositiveTest] {
		test := test
		t.Run(test.Name, func(t *testing.T) {
			if killContext.Err() != nil || (testing.Short() && test.ConfigMinVersion != test.ConfigVersion) {
				t.SkipNow()
			}
			if listSubtests {
				fmt.Println(t.Name())
				return
			}
			t.Parallel()
			err := outer(t, test, false)
			if err != nil {
				t.Error(err)
			}
		})
	}
}

func TestIgnitionBlackBoxNegative(t *testing.T) {
	for _, test := range register.Tests[register.NegativeTest] {
		test := test
		t.Run(test.Name, func(t *testing.T) {
			if killContext.Err() != nil || (testing.Short() && test.ConfigMinVersion != test.ConfigVersion) {
				t.SkipNow()
			}
			if listSubtests {
				fmt.Println(t.Name())
				return
			}
			t.Parallel()
			err := outer(t, test, true)
			if err != nil {
				t.Error(err)
			}
		})
	}
}

func outer(t *testing.T, test types.Test, negativeTests bool) error {
	t.Log(test.Name)

	err := test.ReplaceAllUUIDVars()
	if err != nil {
		return err
	}

	ctx, cancelFunc := context.WithDeadline(killContext, time.Now().Add(testTimeout))
	defer cancelFunc()

	tmpDirectory, err := ioutil.TempDir("/var/tmp", "ignition-blackbox-")
	if err != nil {
		return fmt.Errorf("failed to create a temp dir: %v", err)
	}
	defer os.RemoveAll(tmpDirectory)
	// the tmpDirectory must be 0755 or the tests will fail as the tool will
	// not have permissions to perform some actions in the mounted folders
	err = os.Chmod(tmpDirectory, 0755)
	if err != nil {
		return fmt.Errorf("failed to change mode of temp dir: %v", err)
	}

	systemConfigDir := filepath.Join(tmpDirectory, "system")
	var rootPartition *types.Partition

	// Setup
	err = createFilesFromSlice(systemConfigDir, test.SystemDirFiles)
	// Defer before the error handling because the createFilesFromSlice function
	// can fail after partially-creating things
	defer os.RemoveAll(systemConfigDir)
	if err != nil {
		return err
	}
	for i, disk := range test.In {
		// Set image file path
		disk.ImageFile = filepath.Join(tmpDirectory, fmt.Sprintf("hd%d", i))
		test.Out[i].ImageFile = disk.ImageFile

		// There may be more partitions created by Ignition, so look at the
		// expected output instead of the input to determine image size
		imageSize := test.Out[i].CalculateImageSize()
		if inSize := disk.CalculateImageSize(); inSize > imageSize {
			imageSize = inSize
		}

		// Finish data setup
		for _, part := range disk.Partitions {
			if part.GUID == "" {
				part.GUID = uuid.New()
				if err != nil {
					return err
				}
			}
			err := updateTypeGUID(part)
			if err != nil {
				return err
			}
		}

		disk.SetOffsets()
		for _, part := range test.Out[i].Partitions {
			err := updateTypeGUID(part)
			if err != nil {
				return err
			}
		}
		test.Out[i].SetOffsets()

		if err = setupDisk(ctx, &disk, i, imageSize, tmpDirectory); err != nil {
			return err
		}

		// Creation
		// Move value into the local scope, because disk.ImageFile and device
		// will change by the time this runs
		imageFile := disk.ImageFile
		device := disk.Device
		defer func() {
			if err := os.Remove(imageFile); err != nil {
				t.Errorf("couldn't remove %s: %v", imageFile, err)
			}
		}()
		defer func() {
			if err := destroyDevice(device); err != nil {
				t.Errorf("couldn't destroy device: %v", err)
			}
		}()

		test.Out[i].Device = disk.Device

		err = createFilesForPartitions(ctx, disk.Partitions)
		if err != nil {
			return err
		}

		// Mount device name substitution
		for _, d := range test.MntDevices {
			device := pickPartition(disk.Device, disk.Partitions, d.Label)
			// The device may not be on this disk, if it's not found here let's
			// assume we'll find it on another one and keep going
			if device != "" {
				test.Config = strings.Replace(test.Config, d.Substitution, device, -1)
			}
		}

		// Replace any instance of $disk<num> with the actual loop device
		// that got assigned to it
		test.Config = strings.Replace(test.Config, fmt.Sprintf("$disk%d", i), disk.Device, -1)

		if rootPartition == nil {
			rootPartition = getRootPartition(disk.Partitions)
		}
	}
	if rootPartition == nil {
		return fmt.Errorf("ROOT filesystem not found! A partition labeled ROOT is requred")
	}

	if strings.Contains(test.Config, "passwd") {
		if err := prepareRootPartitionForPasswd(ctx, rootPartition); err != nil {
			return err
		}
	}

	// Validation and cleanup deferral
	for i, disk := range test.Out {
		// Update out structure with mount points & devices
		setExpectedPartitionsDrive(test.In[i].Partitions, disk.Partitions)
	}

	// Let's make sure that all of the devices we needed to substitute names in
	// for were found
	for _, d := range test.MntDevices {
		if strings.Contains(test.Config, d.Substitution) {
			return fmt.Errorf("Didn't find a drive with label: %s", d.Substitution)
		}
	}

	// If we're not expecting the config to be bad, make sure it passes
	// validation.
	if !test.ConfigShouldBeBad {
		_, rpt, err := config.Parse([]byte(test.Config))
		if rpt.IsFatal() {
			return fmt.Errorf("test has bad config: %s", rpt.String())
		}
		if err != nil {
			return fmt.Errorf("error parsing config: %v", err)
		}
	}

	// Ignition config
	if err := ioutil.WriteFile(filepath.Join(tmpDirectory, "config.ign"), []byte(test.Config), 0666); err != nil {
		return fmt.Errorf("error writing config: %v", err)
	}

	// Ignition
	appendEnv := test.Env
	appendEnv = append(appendEnv, "IGNITION_SYSTEM_CONFIG_DIR="+systemConfigDir)

	if !negativeTests {
		if err := runIgnition(t, ctx, "disks", "", tmpDirectory, appendEnv); err != nil {
			return err
		}

		if err := mountPartition(ctx, rootPartition); err != nil {
			return err
		}

		if err := runIgnition(t, ctx, "mount", rootPartition.MountPath, tmpDirectory, appendEnv); err != nil {
			return err
		}

		filesErr := runIgnition(t, ctx, "files", rootPartition.MountPath, tmpDirectory, appendEnv)
		if err := runIgnition(t, ctx, "umount", rootPartition.MountPath, tmpDirectory, appendEnv); err != nil {
			return err
		}
		if err := umountPartition(rootPartition); err != nil {
			return err
		}
		if filesErr != nil {
			return filesErr
		}

		for _, disk := range test.Out {
			err = validateDisk(t, disk)
			if err != nil {
				return err
			}
			err = validateFilesystems(t, disk.Partitions)
			if err != nil {
				return err
			}
			validateFilesDirectoriesAndLinks(t, ctx, disk.Partitions)
		}
		return nil
	} else {
		if err := runIgnition(t, ctx, "disks", "", tmpDirectory, appendEnv); err != nil {
			return nil // error is expected
		}

		if err := mountPartition(ctx, rootPartition); err != nil {
			return err
		}

		if err := runIgnition(t, ctx, "mount", rootPartition.MountPath, tmpDirectory, appendEnv); err != nil {
			return nil // error is expected
		}

		filesErr := runIgnition(t, ctx, "files", rootPartition.MountPath, tmpDirectory, appendEnv)
		if err := runIgnition(t, ctx, "umount", rootPartition.MountPath, tmpDirectory, appendEnv); err != nil {
			return nil
		}
		if err := umountPartition(rootPartition); err != nil {
			return nil
		}
		if filesErr != nil {
			return nil // error is expected
		}
		return fmt.Errorf("Expected failure and ignition succeeded")
	}
}
<commit_msg>tests/blackbox: make runner print config<commit_after>// Copyright 2017 CoreOS, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package blackbox

import (
	"context"
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"syscall"
	"testing"
	"time"

	config "github.com/coreos/ignition/config/v3_0_experimental"
	"github.com/coreos/ignition/tests/register"
	"github.com/coreos/ignition/tests/types"

	// Register the tests
	_ "github.com/coreos/ignition/tests/registry"

	// UUID generation tool
	"github.com/pborman/uuid"
)

var (
	// testTimeout controls how long a given test is allowed to run before being
	// cancelled.
	testTimeout = time.Second * 60
	// somewhat of an abuse of contexts but go's got our hands tied
	killContext = context.TODO()

	// flag for listing all subtests that would be run without running them
	listSubtests = false
)

func TestMain(m *testing.M) {
	interruptChan := make(chan os.Signal, 3)
	signal.Notify(interruptChan, os.Interrupt, syscall.SIGTERM)
	tmp, killCancel := context.WithCancel(context.Background())
	killContext = tmp
	go func() {
		for {
			sig := <-interruptChan
			switch sig {
			case os.Interrupt, syscall.SIGTERM:
				killCancel()
			}
		}
	}()

	flag.BoolVar(&listSubtests, "list", false, "list tests that would be run without running them")
	flag.Parse()

	if !listSubtests {
		httpServer := &HTTPServer{}
		httpServer.Start()
		tftpServer := &TFTPServer{}
		tftpServer.Start()
	}
	os.Exit(m.Run())
}

func TestIgnitionBlackBox(t *testing.T) {
	for _, test := range register.Tests[register.PositiveTest] {
		test := test
		t.Run(test.Name, func(t *testing.T) {
			if killContext.Err() != nil || (testing.Short() && test.ConfigMinVersion != test.ConfigVersion) {
				t.SkipNow()
			}
			if listSubtests {
				fmt.Println(t.Name())
				return
			}
			t.Parallel()
			err := outer(t, test, false)
			if err != nil {
				t.Error(err)
			}
		})
	}
}

func TestIgnitionBlackBoxNegative(t *testing.T) {
	for _, test := range register.Tests[register.NegativeTest] {
		test := test
		t.Run(test.Name, func(t *testing.T) {
			if killContext.Err() != nil || (testing.Short() && test.ConfigMinVersion != test.ConfigVersion) {
				t.SkipNow()
			}
			if listSubtests {
				fmt.Println(t.Name())
				return
			}
			t.Parallel()
			err := outer(t, test, true)
			if err != nil {
				t.Error(err)
			}
		})
	}
}

func outer(t *testing.T, test types.Test, negativeTests bool) error {
	t.Log(test.Name)

	err := test.ReplaceAllUUIDVars()
	if err != nil {
		return err
	}

	ctx, cancelFunc := context.WithDeadline(killContext, time.Now().Add(testTimeout))
	defer cancelFunc()

	tmpDirectory, err := ioutil.TempDir("/var/tmp", "ignition-blackbox-")
	if err != nil {
		return fmt.Errorf("failed to create a temp dir: %v", err)
	}
	defer os.RemoveAll(tmpDirectory)
	// the tmpDirectory must be 0755 or the tests will fail as the tool will
	// not have permissions to perform some actions in the mounted folders
	err = os.Chmod(tmpDirectory, 0755)
	if err != nil {
		return fmt.Errorf("failed to change mode of temp dir: %v", err)
	}

	systemConfigDir := filepath.Join(tmpDirectory, "system")
	var rootPartition *types.Partition

	// Setup
	err = createFilesFromSlice(systemConfigDir, test.SystemDirFiles)
	// Defer before the error handling because the createFilesFromSlice function
	// can fail after partially-creating things
	defer os.RemoveAll(systemConfigDir)
	if err != nil {
		return err
	}
	for i, disk := range test.In {
		// Set image file path
		disk.ImageFile = filepath.Join(tmpDirectory, fmt.Sprintf("hd%d", i))
		test.Out[i].ImageFile = disk.ImageFile

		// There may be more partitions created by Ignition, so look at the
		// expected output instead of the input to determine image size
		imageSize := test.Out[i].CalculateImageSize()
		if inSize := disk.CalculateImageSize(); inSize > imageSize {
			imageSize = inSize
		}

		// Finish data setup
		for _, part := range disk.Partitions {
			if part.GUID == "" {
				part.GUID = uuid.New()
				if err != nil {
					return err
				}
			}
			err := updateTypeGUID(part)
			if err != nil {
				return err
			}
		}

		disk.SetOffsets()
		for _, part := range test.Out[i].Partitions {
			err := updateTypeGUID(part)
			if err != nil {
				return err
			}
		}
		test.Out[i].SetOffsets()

		if err = setupDisk(ctx, &disk, i, imageSize, tmpDirectory); err != nil {
			return err
		}

		// Creation
		// Move value into the local scope, because disk.ImageFile and device
		// will change by the time this runs
		imageFile := disk.ImageFile
		device := disk.Device
		defer func() {
			if err := os.Remove(imageFile); err != nil {
				t.Errorf("couldn't remove %s: %v", imageFile, err)
			}
		}()
		defer func() {
			if err := destroyDevice(device); err != nil {
				t.Errorf("couldn't destroy device: %v", err)
			}
		}()

		test.Out[i].Device = disk.Device

		err = createFilesForPartitions(ctx, disk.Partitions)
		if err != nil {
			return err
		}

		// Mount device name substitution
		for _, d := range test.MntDevices {
			device := pickPartition(disk.Device, disk.Partitions, d.Label)
			// The device may not be on this disk, if it's not found here let's
			// assume we'll find it on another one and keep going
			if device != "" {
				test.Config = strings.Replace(test.Config, d.Substitution, device, -1)
			}
		}

		// Replace any instance of $disk<num> with the actual loop device
		// that got assigned to it
		test.Config = strings.Replace(test.Config, fmt.Sprintf("$disk%d", i), disk.Device, -1)

		if rootPartition == nil {
			rootPartition = getRootPartition(disk.Partitions)
		}
	}
	if rootPartition == nil {
		return fmt.Errorf("ROOT filesystem not found! A partition labeled ROOT is requred")
	}

	if strings.Contains(test.Config, "passwd") {
		if err := prepareRootPartitionForPasswd(ctx, rootPartition); err != nil {
			return err
		}
	}

	// Validation and cleanup deferral
	for i, disk := range test.Out {
		// Update out structure with mount points & devices
		setExpectedPartitionsDrive(test.In[i].Partitions, disk.Partitions)
	}

	// Let's make sure that all of the devices we needed to substitute names in
	// for were found
	for _, d := range test.MntDevices {
		if strings.Contains(test.Config, d.Substitution) {
			return fmt.Errorf("Didn't find a drive with label: %s", d.Substitution)
		}
	}

	t.Logf("Rendered Ignition Config:\n%s", test.Config)

	// If we're not expecting the config to be bad, make sure it passes
	// validation.
	if !test.ConfigShouldBeBad {
		_, rpt, err := config.Parse([]byte(test.Config))
		if rpt.IsFatal() {
			return fmt.Errorf("test has bad config: %s", rpt.String())
		}
		if err != nil {
			return fmt.Errorf("error parsing config: %v", err)
		}
	}

	// Ignition config
	if err := ioutil.WriteFile(filepath.Join(tmpDirectory, "config.ign"), []byte(test.Config), 0666); err != nil {
		return fmt.Errorf("error writing config: %v", err)
	}

	// Ignition
	appendEnv := test.Env
	appendEnv = append(appendEnv, "IGNITION_SYSTEM_CONFIG_DIR="+systemConfigDir)

	if !negativeTests {
		if err := runIgnition(t, ctx, "disks", "", tmpDirectory, appendEnv); err != nil {
			return err
		}

		if err := mountPartition(ctx, rootPartition); err != nil {
			return err
		}

		if err := runIgnition(t, ctx, "mount", rootPartition.MountPath, tmpDirectory, appendEnv); err != nil {
			return err
		}

		filesErr := runIgnition(t, ctx, "files", rootPartition.MountPath, tmpDirectory, appendEnv)
		if err := runIgnition(t, ctx, "umount", rootPartition.MountPath, tmpDirectory, appendEnv); err != nil {
			return err
		}
		if err := umountPartition(rootPartition); err != nil {
			return err
		}
		if filesErr != nil {
			return filesErr
		}

		for _, disk := range test.Out {
			err = validateDisk(t, disk)
			if err != nil {
				return err
			}
			err = validateFilesystems(t, disk.Partitions)
			if err != nil {
				return err
			}
			validateFilesDirectoriesAndLinks(t, ctx, disk.Partitions)
		}
		return nil
	} else {
		if err := runIgnition(t, ctx, "disks", "", tmpDirectory, appendEnv); err != nil {
			return nil // error is expected
		}

		if err := mountPartition(ctx, rootPartition); err != nil {
			return err
		}

		if err := runIgnition(t, ctx, "mount", rootPartition.MountPath, tmpDirectory, appendEnv); err != nil {
			return nil // error is expected
		}

		filesErr := runIgnition(t, ctx, "files", rootPartition.MountPath, tmpDirectory, appendEnv)
		if err := runIgnition(t, ctx, "umount", rootPartition.MountPath, tmpDirectory, appendEnv); err != nil {
			return nil
		}
		if err := umountPartition(rootPartition); err != nil {
			return nil
		}
		if filesErr != nil {
			return nil // error is expected
		}
		return fmt.Errorf("Expected failure and ignition succeeded")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"github.com/lib/pq"
)

func groupGoods() {
	// load image and urls infos
	hashToUrlIds := make(map[string][]int64)
	rows, err := db.Query(`SELECT url_id, encode(sha512_16k, 'base64') FROM urls
		WHERE sha512_16k IS NOT NULL
		--ORDER BY url_id DESC
		--LIMIT 50000 -- DEBUG
		`)
	ce(err, "query")
	n := 0
	for rows.Next() {
		var urlId int64
		var hash string
		ce(rows.Scan(&urlId, &hash), "scan")
		hashToUrlIds[hash] = append(hashToUrlIds[hash], urlId)
		n++
		if n%10000 == 0 {
			pt("%d\n", n)
		}
	}
	ce(rows.Err(), "rows err")
	pt("urls loaded\n")

	urlIdToGoodIds := make(map[int64][]int64)
	rows, err = db.Query(`SELECT good_id, url_id FROM images
		--ORDER BY image_id DESC
		--LIMIT 50000 -- DEBUG
		`)
	ce(err, "query")
	n = 0
	for rows.Next() {
		var goodId, urlId int64
		ce(rows.Scan(&goodId, &urlId), "scan")
		urlIdToGoodIds[urlId] = append(urlIdToGoodIds[urlId], goodId)
		n++
		if n%10000 == 0 {
			pt("%d\n", n)
		}
	}
	ce(rows.Err(), "rows err")
	pt("images loaded\n")

	txCount := 0
	tx := db.MustBegin()

check:

	// get good id
	var goodId int64
	ce(tx.Get(&goodId, `SELECT
		good_id
		FROM goods
		WHERE group_id IS NULL
		AND images_collected = true
		ORDER BY good_id DESC
		LIMIT 1`), "get good id")

	// get good hashes
	var hashes pq.StringArray
	ce(tx.Select(&hashes, `SELECT
		encode(sha512_16k, 'base64')
		FROM images i
		LEFT JOIN urls USING(url_id)
		WHERE
		i.good_id = $1
		`,
		goodId,
	), "select hashes")

	// stat
	matches := make(map[int64]int)
	for _, hash := range hashes {
		urlIds := hashToUrlIds[hash]
		for _, urlId := range urlIds {
			for _, rightId := range urlIdToGoodIds[urlId] {
				matches[rightId]++
			}
		}
	}
	has := false
	for rightId, n := range matches {
		if n >= 10 {
			pt("-> %d %d\n", goodId, rightId)
			// 不到的话就算了吧
			has = true
			_, err := tx.Exec(`UPDATE goods
				SET group_id = $1
				WHERE good_id = $2
				`,
				goodId,
				rightId,
			)
			ce(err, "update goods")
		}
	}
	if has {
		_, err := tx.Exec(`UPDATE goods
			SET group_id = $1
			WHERE good_id = $1
			`,
			goodId,
		)
		ce(err, "update goods")
	} else {
		_, err := tx.Exec(`UPDATE goods
				SET group_id = -1
				WHERE good_id = $1
				`,
			goodId,
		)
		ce(err, "update goods")
	}

	txCount++
	if txCount >= 256 {
		ce(tx.Commit(), "commit")
		tx = db.MustBegin()
		txCount = 0
	}

	goto check
}
<commit_msg>use arrays<commit_after>package main

type Hash [64]byte

func groupGoods() {
	// load image and urls infos
	hashToUrlIds := make(map[Hash][]int64)
	rows, err := db.Query(`SELECT url_id, sha512_16k FROM urls
		WHERE sha512_16k IS NOT NULL
		--ORDER BY url_id DESC
		--LIMIT 50000 -- DEBUG
		`)
	ce(err, "query")
	n := 0
	for rows.Next() {
		var urlId int64
		var h []byte
		ce(rows.Scan(&urlId, &h), "scan")
		var hash Hash
		copy(hash[:], h[:])
		hashToUrlIds[hash] = append(hashToUrlIds[hash], urlId)
		n++
		if n%10000 == 0 {
			pt("%d\n", n)
		}
	}
	ce(rows.Err(), "rows err")
	pt("urls loaded\n")

	urlIdToGoodIds := make(map[int64][]int64)
	rows, err = db.Query(`SELECT good_id, url_id FROM images
		--ORDER BY image_id DESC
		--LIMIT 50000 -- DEBUG
		`)
	ce(err, "query")
	n = 0
	for rows.Next() {
		var goodId, urlId int64
		ce(rows.Scan(&goodId, &urlId), "scan")
		urlIdToGoodIds[urlId] = append(urlIdToGoodIds[urlId], goodId)
		n++
		if n%10000 == 0 {
			pt("%d\n", n)
		}
	}
	ce(rows.Err(), "rows err")
	pt("images loaded\n")

	txCount := 0
	tx := db.MustBegin()

check:

	// get good id
	var goodId int64
	ce(tx.Get(&goodId, `SELECT
		good_id
		FROM goods
		WHERE group_id IS NULL
		AND images_collected = true
		ORDER BY good_id DESC
		LIMIT 1`), "get good id")

	// get good hashes
	var hs [][]byte
	ce(tx.Select(&hs, `SELECT
		sha512_16k
		FROM images i
		LEFT JOIN urls USING(url_id)
		WHERE
		i.good_id = $1
		`,
		goodId,
	), "select hashes")
	var hashes []Hash
	for _, h := range hs {
		var hash Hash
		copy(hash[:], h[:])
		hashes = append(hashes, hash)
	}

	// stat
	matches := make(map[int64]int)
	for _, hash := range hashes {
		urlIds := hashToUrlIds[hash]
		for _, urlId := range urlIds {
			for _, rightId := range urlIdToGoodIds[urlId] {
				matches[rightId]++
			}
		}
	}
	has := false
	for rightId, n := range matches {
		if n >= 10 {
			pt("-> %d %d\n", goodId, rightId)
			// 不到的话就算了吧
			has = true
			_, err := tx.Exec(`UPDATE goods
				SET group_id = $1
				WHERE good_id = $2
				`,
				goodId,
				rightId,
			)
			ce(err, "update goods")
		}
	}
	if has {
		_, err := tx.Exec(`UPDATE goods
			SET group_id = $1
			WHERE good_id = $1
			`,
			goodId,
		)
		ce(err, "update goods")
	} else {
		_, err := tx.Exec(`UPDATE goods
				SET group_id = -1
				WHERE good_id = $1
				`,
			goodId,
		)
		ce(err, "update goods")
	}

	txCount++
	if txCount >= 256 {
		ce(tx.Commit(), "commit")
		tx = db.MustBegin()
		txCount = 0
	}

	goto check
}
<|endoftext|>


<commit_before>package main

import "fmt"

func main() {
  fmt.Printf("Hello, world.\n")
}

<commit_msg>add comment<commit_after>package main

import "fmt"

func main() {
  fmt.Printf("Hello, world.\n")
}

// test2
<|endoftext|>


<commit_before>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package context

import (
	"github.com/juju/cmd"
	"github.com/juju/errors"

	"github.com/juju/juju/workload"
)

// StatusSetCmdName is the name of the payload status-set command.
const StatusSetCmdName = "payload-status-set"

// NewStatusSetCmd returns a new StatusSetCmd that wraps the given context.
func NewStatusSetCmd(ctx HookContext) (*StatusSetCmd, error) {
	compCtx, err := ContextComponent(ctx)
	if err != nil {
		// The component wasn't tracked properly.
		return nil, errors.Trace(err)
	}
	return &StatusSetCmd{api: compCtx}, nil
}

// StatusSetCmd is a command that registers a payload with juju.
type StatusSetCmd struct {
	cmd.CommandBase

	api    Component
	class  string
	id     string
	status string
}

// Info implements cmd.Command.
func (c StatusSetCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:    StatusSetCmdName,
		Args:    "<class> <id> <status>",
		Purpose: "update the status of a payload",
		Doc: `
"payload-status-set" is used while a hook (update-status) is running to update the
current status of a registered payload. The <class> and <id> provided must match a
payload that has been previously registered with juju using payload-register.
The <status> must be on of the follow: starting, started, stopping, stopped
`,
	}
}

// Init implements cmd.Command.
func (c *StatusSetCmd) Init(args []string) error {
	if len(args) < 3 {
		return errors.Errorf("missing required arguments")
	}
	c.class = args[0]
	c.id = args[1]
	c.status = args[2]
	return nil
}

// Run implements cmd.Command.
func (c *StatusSetCmd) Run(ctx *cmd.Context) error {
	if err := c.validate(ctx); err != nil {
		return errors.Trace(err)
	}

	if err := c.api.SetStatus(c.class, c.id, c.status); err != nil {
		return errors.Trace(err)
	}

	// We flush to state immedeiately so that status reflects the
	// workload correctly.
	if err := c.api.Flush(); err != nil {
		return errors.Trace(err)
	}

	return nil
}

func (c *StatusSetCmd) validate(ctx *cmd.Context) error {
	return workload.ValidateState(c.status)
}
<commit_msg>check empty args<commit_after>// Copyright 2015 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package context

import (
	"github.com/juju/cmd"
	"github.com/juju/errors"

	"github.com/juju/juju/workload"
)

// StatusSetCmdName is the name of the payload status-set command.
const StatusSetCmdName = "payload-status-set"

// NewStatusSetCmd returns a new StatusSetCmd that wraps the given context.
func NewStatusSetCmd(ctx HookContext) (*StatusSetCmd, error) {
	compCtx, err := ContextComponent(ctx)
	if err != nil {
		// The component wasn't tracked properly.
		return nil, errors.Trace(err)
	}
	return &StatusSetCmd{api: compCtx}, nil
}

// StatusSetCmd is a command that registers a payload with juju.
type StatusSetCmd struct {
	cmd.CommandBase

	api    Component
	class  string
	id     string
	status string
}

// Info implements cmd.Command.
func (c StatusSetCmd) Info() *cmd.Info {
	return &cmd.Info{
		Name:    StatusSetCmdName,
		Args:    "<class> <id> <status>",
		Purpose: "update the status of a payload",
		Doc: `
"payload-status-set" is used while a hook (update-status) is running to update the
current status of a registered payload. The <class> and <id> provided must match a
payload that has been previously registered with juju using payload-register.
The <status> must be on of the follow: starting, started, stopping, stopped
`,
	}
}

// Init implements cmd.Command.
func (c *StatusSetCmd) Init(args []string) error {
	if len(args) < 3 {
		return errors.Errorf("missing required arguments")
	}
	c.class = args[0]
	c.id = args[1]
	c.status = args[2]
	return cmd.CheckEmpty(args[3:])
}

// Run implements cmd.Command.
func (c *StatusSetCmd) Run(ctx *cmd.Context) error {
	if err := c.validate(ctx); err != nil {
		return errors.Trace(err)
	}

	if err := c.api.SetStatus(c.class, c.id, c.status); err != nil {
		return errors.Trace(err)
	}

	// We flush to state immedeiately so that status reflects the
	// workload correctly.
	if err := c.api.Flush(); err != nil {
		return errors.Trace(err)
	}

	return nil
}

func (c *StatusSetCmd) validate(ctx *cmd.Context) error {
	return workload.ValidateState(c.status)
}
<|endoftext|>


<commit_before>// Copyright 2019-2020 Google LLC
//
// 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 main

import (
	"context"
	"log"
	"strings"
)

// In-memory implementation of StorageService, for local testing
type inMemStore struct {
	teams    map[string]Team
	periods  map[string]map[string]Period
	settings Settings
}

func makeInMemStore() StorageService {
	teams := map[string]Team{
		"team1": {ID: "team1", DisplayName: "Team team1"},
		"team2": {ID: "team2", DisplayName: "Team team2"},
	}
	periods := map[string]map[string]Period{
		"team1": {
			"2018q4": makeFakePeriod("2018q4"),
			"2019q1": makeFakePeriod("2019q1"),
		},
		"team2": {
			"2018q4": makeFakePeriod("2018q4"),
		},
	}
	settings := Settings{
		ImproveURL: "https://github.com/google/peoplemath",
	}
	return &inMemStore{teams: teams, periods: periods, settings: settings}
}

func (s *inMemStore) GetAllTeams(ctx context.Context) ([]Team, error) {
	teamsSlice := make([]Team, 0, len(s.teams))
	for _, t := range s.teams {
		teamsSlice = append(teamsSlice, t)
	}
	return teamsSlice, nil
}

func (s *inMemStore) GetTeam(ctx context.Context, teamID string) (Team, bool, error) {
	team, ok := s.teams[teamID]
	return team, ok, nil
}

func (s *inMemStore) CreateTeam(ctx context.Context, team Team) error {
	s.teams[team.ID] = team
	s.periods[team.ID] = map[string]Period{}
	log.Printf("Added new team %s", team.ID)
	return nil
}

func (s *inMemStore) UpdateTeam(ctx context.Context, team Team) error {
	s.teams[team.ID] = team
	log.Printf("Updated team %s", team.ID)
	return nil
}

func (s *inMemStore) GetAllPeriods(ctx context.Context, teamID string) ([]Period, bool, error) {
	if periodsByName, ok := s.periods[teamID]; ok {
		periodSlice := make([]Period, 0, len(periodsByName))
		for _, p := range periodsByName {
			periodSlice = append(periodSlice, p)
		}
		return periodSlice, true, nil
	}
	return []Period{}, false, nil
}

func (s *inMemStore) GetPeriod(ctx context.Context, teamID, periodID string) (Period, bool, error) {
	if periodsByName, ok := s.periods[teamID]; ok {
		if period, ok := periodsByName[periodID]; ok {
			return period, true, nil
		}
	}
	return Period{}, false, nil
}

func (s *inMemStore) CreatePeriod(ctx context.Context, teamID string, period Period) error {
	if periodsByName, ok := s.periods[teamID]; ok {
		periodsByName[period.ID] = period
		log.Printf("Added period '%s' for team '%s': %v", period.ID, teamID, period)
	}
	return nil
}

func (s *inMemStore) UpdatePeriod(ctx context.Context, teamID string, period Period) error {
	if periodsByName, ok := s.periods[teamID]; ok {
		periodsByName[period.ID] = period
		log.Printf("Updated period '%s' for team '%s': %v", period.ID, teamID, period)
	}
	return nil
}

func (s *inMemStore) GetSettings(ctx context.Context) (Settings, error) {
	return s.settings, nil
}

func (s *inMemStore) Close() error {
	return nil
}

func makeFakePeriod(id string) Period {
	buckets := []Bucket{
		Bucket{
			DisplayName:          "First bucket",
			AllocationPercentage: 40,
			Objectives: []Objective{
				Objective{
					Name:             "First objective",
					ResourceEstimate: 10,
					CommitmentType:   "Committed",
					Assignments: []Assignment{
						Assignment{
							PersonID:   "alice",
							Commitment: 5,
						},
						Assignment{
							PersonID:   "bob",
							Commitment: 5,
						},
					},
					Groups: []ObjectiveGroup{
						ObjectiveGroup{
							GroupType: "Project",
							GroupName: "Project 1",
						},
					},
					Tags: []ObjectiveTag{},
				},
				Objective{
					Name:             "Second objective",
					ResourceEstimate: 15,
					CommitmentType:   "Aspirational",
					Notes:            "Some notes",
					Assignments: []Assignment{
						Assignment{
							PersonID:   "bob",
							Commitment: 2,
						},
					},
					Groups: []ObjectiveGroup{
						ObjectiveGroup{
							GroupType: "Project",
							GroupName: "Project 2",
						},
					},
					Tags: []ObjectiveTag{},
				},
			},
		},
		Bucket{
			DisplayName:          "Second bucket",
			AllocationPercentage: 40,
			Objectives: []Objective{
				Objective{
					Name:             "Third objective",
					ResourceEstimate: 2,
					CommitmentType:   "Aspirational",
					Assignments:      []Assignment{},
					Groups: []ObjectiveGroup{
						ObjectiveGroup{
							GroupType: "Project",
							GroupName: "Project 1",
						},
					},
					Tags: []ObjectiveTag{},
				},
				Objective{
					Name:             "Fourth objective",
					ResourceEstimate: 8,
					CommitmentType:   "Aspirational",
					Assignments: []Assignment{
						Assignment{
							PersonID:   "charlie",
							Commitment: 8,
						},
					},
					Groups: []ObjectiveGroup{},
					Tags:   []ObjectiveTag{},
				},
			},
		},
		Bucket{
			DisplayName:          "Third bucket",
			AllocationPercentage: 20,
			Objectives:           []Objective{},
		},
	}
	people := []Person{
		Person{
			ID:           "alice",
			DisplayName:  "Alice Atkins",
			Availability: 5,
		},
		Person{
			ID:           "bob",
			DisplayName:  "Bob Brewster",
			Availability: 7,
		},
		Person{
			ID:           "charlie",
			DisplayName:  "Charlie Case",
			Availability: 8,
		},
	}
	return Period{
		ID:                     id,
		DisplayName:            strings.ToUpper(id),
		Unit:                   "person weeks",
		NotesURL:               "https://github.com/google/peoplemath",
		MaxCommittedPercentage: 50,
		Buckets:                buckets,
		People:                 people,
	}
}
<commit_msg>Add tags to inmemstore start state (#21)<commit_after>// Copyright 2019-2020 Google LLC
//
// 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 main

import (
	"context"
	"log"
	"strings"
)

// In-memory implementation of StorageService, for local testing
type inMemStore struct {
	teams    map[string]Team
	periods  map[string]map[string]Period
	settings Settings
}

func makeInMemStore() StorageService {
	teams := map[string]Team{
		"team1": {ID: "team1", DisplayName: "Team team1"},
		"team2": {ID: "team2", DisplayName: "Team team2"},
	}
	periods := map[string]map[string]Period{
		"team1": {
			"2018q4": makeFakePeriod("2018q4"),
			"2019q1": makeFakePeriod("2019q1"),
		},
		"team2": {
			"2018q4": makeFakePeriod("2018q4"),
		},
	}
	settings := Settings{
		ImproveURL: "https://github.com/google/peoplemath",
	}
	return &inMemStore{teams: teams, periods: periods, settings: settings}
}

func (s *inMemStore) GetAllTeams(ctx context.Context) ([]Team, error) {
	teamsSlice := make([]Team, 0, len(s.teams))
	for _, t := range s.teams {
		teamsSlice = append(teamsSlice, t)
	}
	return teamsSlice, nil
}

func (s *inMemStore) GetTeam(ctx context.Context, teamID string) (Team, bool, error) {
	team, ok := s.teams[teamID]
	return team, ok, nil
}

func (s *inMemStore) CreateTeam(ctx context.Context, team Team) error {
	s.teams[team.ID] = team
	s.periods[team.ID] = map[string]Period{}
	log.Printf("Added new team %s", team.ID)
	return nil
}

func (s *inMemStore) UpdateTeam(ctx context.Context, team Team) error {
	s.teams[team.ID] = team
	log.Printf("Updated team %s", team.ID)
	return nil
}

func (s *inMemStore) GetAllPeriods(ctx context.Context, teamID string) ([]Period, bool, error) {
	if periodsByName, ok := s.periods[teamID]; ok {
		periodSlice := make([]Period, 0, len(periodsByName))
		for _, p := range periodsByName {
			periodSlice = append(periodSlice, p)
		}
		return periodSlice, true, nil
	}
	return []Period{}, false, nil
}

func (s *inMemStore) GetPeriod(ctx context.Context, teamID, periodID string) (Period, bool, error) {
	if periodsByName, ok := s.periods[teamID]; ok {
		if period, ok := periodsByName[periodID]; ok {
			return period, true, nil
		}
	}
	return Period{}, false, nil
}

func (s *inMemStore) CreatePeriod(ctx context.Context, teamID string, period Period) error {
	if periodsByName, ok := s.periods[teamID]; ok {
		periodsByName[period.ID] = period
		log.Printf("Added period '%s' for team '%s': %v", period.ID, teamID, period)
	}
	return nil
}

func (s *inMemStore) UpdatePeriod(ctx context.Context, teamID string, period Period) error {
	if periodsByName, ok := s.periods[teamID]; ok {
		periodsByName[period.ID] = period
		log.Printf("Updated period '%s' for team '%s': %v", period.ID, teamID, period)
	}
	return nil
}

func (s *inMemStore) GetSettings(ctx context.Context) (Settings, error) {
	return s.settings, nil
}

func (s *inMemStore) Close() error {
	return nil
}

func makeFakePeriod(id string) Period {
	buckets := []Bucket{
		Bucket{
			DisplayName:          "First bucket",
			AllocationPercentage: 40,
			Objectives: []Objective{
				Objective{
					Name:             "First objective",
					ResourceEstimate: 10,
					CommitmentType:   "Committed",
					Assignments: []Assignment{
						Assignment{
							PersonID:   "alice",
							Commitment: 5,
						},
						Assignment{
							PersonID:   "bob",
							Commitment: 5,
						},
					},
					Groups: []ObjectiveGroup{
						ObjectiveGroup{
							GroupType: "Project",
							GroupName: "Project 1",
						},
					},
					Tags: []ObjectiveTag{},
				},
				Objective{
					Name:             "Second objective",
					ResourceEstimate: 15,
					CommitmentType:   "Aspirational",
					Notes:            "Some notes",
					Assignments: []Assignment{
						Assignment{
							PersonID:   "bob",
							Commitment: 2,
						},
					},
					Groups: []ObjectiveGroup{
						ObjectiveGroup{
							GroupType: "Project",
							GroupName: "Project 2",
						},
					},
					Tags: []ObjectiveTag{
						ObjectiveTag{
							Name: "tag1",
						},
						ObjectiveTag{
							Name: "tag2",
						},
					},
				},
			},
		},
		Bucket{
			DisplayName:          "Second bucket",
			AllocationPercentage: 40,
			Objectives: []Objective{
				Objective{
					Name:             "Third objective",
					ResourceEstimate: 2,
					CommitmentType:   "Aspirational",
					Assignments:      []Assignment{},
					Groups: []ObjectiveGroup{
						ObjectiveGroup{
							GroupType: "Project",
							GroupName: "Project 1",
						},
					},
					Tags: []ObjectiveTag{},
				},
				Objective{
					Name:             "Fourth objective",
					ResourceEstimate: 8,
					CommitmentType:   "Aspirational",
					Assignments: []Assignment{
						Assignment{
							PersonID:   "charlie",
							Commitment: 8,
						},
					},
					Groups: []ObjectiveGroup{},
					Tags:   []ObjectiveTag{},
				},
			},
		},
		Bucket{
			DisplayName:          "Third bucket",
			AllocationPercentage: 20,
			Objectives:           []Objective{},
		},
	}
	people := []Person{
		Person{
			ID:           "alice",
			DisplayName:  "Alice Atkins",
			Availability: 5,
		},
		Person{
			ID:           "bob",
			DisplayName:  "Bob Brewster",
			Availability: 7,
		},
		Person{
			ID:           "charlie",
			DisplayName:  "Charlie Case",
			Availability: 8,
		},
	}
	return Period{
		ID:                     id,
		DisplayName:            strings.ToUpper(id),
		Unit:                   "person weeks",
		NotesURL:               "https://github.com/google/peoplemath",
		MaxCommittedPercentage: 50,
		Buckets:                buckets,
		People:                 people,
	}
}
<|endoftext|>


<commit_before>package maimok

import (
	"bytes"
	"crypto/rand"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"path/filepath"

	libvirt "github.com/libvirt/libvirt-go"
	"github.com/satori/go.uuid"
)

// VM data model
type VM struct {
	ID     uint    `json:"id"`
	Name   string  `json:"name"`
	Memory uint64  `json:"memory"`
	State  VMState `json:"state"`
}

// VMState is the state of the vm
type VMState int

const (
	// Running vm
	Running VMState = iota + 1
	// Stopped vm
	Stopped
)

// CreateVMStruct struct for the CreateVM call
type CreateVMStruct struct {
	// Generated
	ID         string
	MACAddress string
	// Provided by API
	Name        string
	Hostname    string
	RAMMB       uint
	DiskSpaceGB uint
	Image       string
	IPAddress   string
	// Set in config file
	SSHKey  string
	Gateway string
	Netmask string
}

func generateMACAddress() (string, error) {
	buf := make(net.HardwareAddr, 6)
	_, err := rand.Read(buf)
	if err != nil {
		return "", err
	}
	buf[0] = (buf[0] | 2) & 0xfe // Set local bit, ensure unicast address
	return buf.String(), nil
}

// CreateVM createds a virtual machine
func CreateVM(state *globalState, createVM CreateVMStruct) error {
	createVM.ID = uuid.NewV4().String()
	createVM.SSHKey = state.config.SSHKey
	createVM.Gateway = state.config.Gateway
	createVM.Netmask = state.config.Netmask

	macAddress, err := generateMACAddress()
	if err != nil {
		return fmt.Errorf("Cannot generate MAC Address, %s", err)
	}
	createVM.MACAddress = macAddress

	// create config iso
	dir, err := ioutil.TempDir("", "maimok")
	if err != nil {
		return fmt.Errorf("Cannot create temp directory, %s", err)
	}
	defer os.RemoveAll(dir)

	metaDataFile, err := os.OpenFile(filepath.Join(dir, "meta-data"), os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		return err
	}
	defer metaDataFile.Close()
	if err := state.tpl.ExecuteTemplate(metaDataFile, "meta-data.yml", createVM); err != nil {
		return err
	}

	userDataFile, err := os.OpenFile(filepath.Join(dir, "user-data"), os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		return err
	}
	defer userDataFile.Close()
	if err := state.tpl.ExecuteTemplate(userDataFile, "user-data.yml", createVM); err != nil {
		return err
	}

	networkConfigFile, err := os.OpenFile(filepath.Join(dir, "network-config"), os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		return err
	}
	defer networkConfigFile.Close()
	if err := state.tpl.ExecuteTemplate(networkConfigFile, "network-config.yml", createVM); err != nil {
		return err
	}

	cmd := exec.Command("genisoimage", "-volid", "cidata", "-joliet", "-rock",
		metaDataFile.Name(), userDataFile.Name(), networkConfigFile.Name())
	var isoFile bytes.Buffer
	cmd.Stdout = &isoFile
	err = cmd.Run()
	if err != nil {
		return fmt.Errorf("Generation of iso image failed: %s", err)
	}

	storagePool, err := state.conn.LookupStoragePoolByName("default")
	if err != nil {
		return err
	}

	// create iso volume
	buf := new(bytes.Buffer)
	if err := state.tpl.ExecuteTemplate(buf, "iso-volume.xml", createVM); err != nil {
		return err
	}
	volume, err := storagePool.StorageVolCreateXML(buf.String(), 0)
	if err != nil {
		return err
	}

	stream, err := state.conn.NewStream(0)
	if err != nil {
		return err
	}
	volume.Upload(stream, 0, 0, 0)
	len, err := stream.Send(isoFile.Bytes())
	if err != nil {
		return fmt.Errorf("Could not send config iso image to storage pool: %s", err)
	}
	if len < buf.Len() {
		return fmt.Errorf("Could not send all iso image data to storage pool")
	}
	if err := stream.Finish(); err != nil {
		return fmt.Errorf("Could not send config iso image to storage pool: %s", err)
	}

	// create harddisk volume
	buf = new(bytes.Buffer)
	if err := state.tpl.ExecuteTemplate(buf, "volume.xml", createVM); err != nil {
		return err
	}
	if _, err = storagePool.StorageVolCreateXML(buf.String(), 0); err != nil {
		return err
	}

	// define domain
	buf = new(bytes.Buffer)
	if err = state.tpl.ExecuteTemplate(buf, "domain.xml", createVM); err != nil {
		return err
	}
	domain, err := state.conn.DomainDefineXML(buf.String())
	if err != nil {
		return fmt.Errorf("Failed to define domain: %s", err)
	}

	// start domain
	if err = domain.Create(); err != nil {
		return fmt.Errorf("Failed to start domain: %s", err)
	}

	return nil
}

// ListVMs returns a list of all virtual machines
func ListVMs(state *globalState) []*VM {
	domains, err := state.conn.ListAllDomains(libvirt.CONNECT_LIST_DOMAINS_ACTIVE)
	if err != nil {
		return nil
	}

	vms := []*VM{}
	for _, domain := range domains {
		name, _ := domain.GetName()
		memory, _ := domain.GetMaxMemory()
		id, _ := domain.GetID()
		state, _, _ := domain.GetState()

		var vmState VMState
		if state == libvirt.DOMAIN_RUNNING {
			vmState = Running
		} else {
			vmState = Stopped
		}

		vms = append(vms, &VM{ID: id, Name: name, Memory: memory, State: vmState})
	}
	return vms
}
<commit_msg>backend: list all vms, not just running<commit_after>package maimok

import (
	"bytes"
	"crypto/rand"
	"fmt"
	"io/ioutil"
	"net"
	"os"
	"os/exec"
	"path/filepath"

	libvirt "github.com/libvirt/libvirt-go"
	"github.com/satori/go.uuid"
)

// VM data model
type VM struct {
	ID     uint    `json:"id"`
	Name   string  `json:"name"`
	Memory uint64  `json:"memory"`
	State  VMState `json:"state"`
}

// VMState is the state of the vm
type VMState int

const (
	// Running vm
	Running VMState = iota + 1
	// Stopped vm
	Stopped
)

// CreateVMStruct struct for the CreateVM call
type CreateVMStruct struct {
	// Generated
	ID         string
	MACAddress string
	// Provided by API
	Name        string
	Hostname    string
	RAMMB       uint
	DiskSpaceGB uint
	Image       string
	IPAddress   string
	// Set in config file
	SSHKey  string
	Gateway string
	Netmask string
}

func generateMACAddress() (string, error) {
	buf := make(net.HardwareAddr, 6)
	_, err := rand.Read(buf)
	if err != nil {
		return "", err
	}
	buf[0] = (buf[0] | 2) & 0xfe // Set local bit, ensure unicast address
	return buf.String(), nil
}

// CreateVM createds a virtual machine
func CreateVM(state *globalState, createVM CreateVMStruct) error {
	createVM.ID = uuid.NewV4().String()
	createVM.SSHKey = state.config.SSHKey
	createVM.Gateway = state.config.Gateway
	createVM.Netmask = state.config.Netmask

	macAddress, err := generateMACAddress()
	if err != nil {
		return fmt.Errorf("Cannot generate MAC Address, %s", err)
	}
	createVM.MACAddress = macAddress

	// create config iso
	dir, err := ioutil.TempDir("", "maimok")
	if err != nil {
		return fmt.Errorf("Cannot create temp directory, %s", err)
	}
	defer os.RemoveAll(dir)

	metaDataFile, err := os.OpenFile(filepath.Join(dir, "meta-data"), os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		return err
	}
	defer metaDataFile.Close()
	if err := state.tpl.ExecuteTemplate(metaDataFile, "meta-data.yml", createVM); err != nil {
		return err
	}

	userDataFile, err := os.OpenFile(filepath.Join(dir, "user-data"), os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		return err
	}
	defer userDataFile.Close()
	if err := state.tpl.ExecuteTemplate(userDataFile, "user-data.yml", createVM); err != nil {
		return err
	}

	networkConfigFile, err := os.OpenFile(filepath.Join(dir, "network-config"), os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		return err
	}
	defer networkConfigFile.Close()
	if err := state.tpl.ExecuteTemplate(networkConfigFile, "network-config.yml", createVM); err != nil {
		return err
	}

	cmd := exec.Command("genisoimage", "-volid", "cidata", "-joliet", "-rock",
		metaDataFile.Name(), userDataFile.Name(), networkConfigFile.Name())
	var isoFile bytes.Buffer
	cmd.Stdout = &isoFile
	err = cmd.Run()
	if err != nil {
		return fmt.Errorf("Generation of iso image failed: %s", err)
	}

	storagePool, err := state.conn.LookupStoragePoolByName("default")
	if err != nil {
		return err
	}

	// create iso volume
	buf := new(bytes.Buffer)
	if err := state.tpl.ExecuteTemplate(buf, "iso-volume.xml", createVM); err != nil {
		return err
	}
	volume, err := storagePool.StorageVolCreateXML(buf.String(), 0)
	if err != nil {
		return err
	}

	stream, err := state.conn.NewStream(0)
	if err != nil {
		return err
	}
	volume.Upload(stream, 0, 0, 0)
	len, err := stream.Send(isoFile.Bytes())
	if err != nil {
		return fmt.Errorf("Could not send config iso image to storage pool: %s", err)
	}
	if len < buf.Len() {
		return fmt.Errorf("Could not send all iso image data to storage pool")
	}
	if err := stream.Finish(); err != nil {
		return fmt.Errorf("Could not send config iso image to storage pool: %s", err)
	}

	// create harddisk volume
	buf = new(bytes.Buffer)
	if err := state.tpl.ExecuteTemplate(buf, "volume.xml", createVM); err != nil {
		return err
	}
	if _, err = storagePool.StorageVolCreateXML(buf.String(), 0); err != nil {
		return err
	}

	// define domain
	buf = new(bytes.Buffer)
	if err = state.tpl.ExecuteTemplate(buf, "domain.xml", createVM); err != nil {
		return err
	}
	domain, err := state.conn.DomainDefineXML(buf.String())
	if err != nil {
		return fmt.Errorf("Failed to define domain: %s", err)
	}

	// start domain
	if err = domain.Create(); err != nil {
		return fmt.Errorf("Failed to start domain: %s", err)
	}

	return nil
}

// ListVMs returns a list of all virtual machines
func ListVMs(state *globalState) []*VM {
	domains, err := state.conn.ListAllDomains(0)
	if err != nil {
		return nil
	}

	vms := []*VM{}
	for _, domain := range domains {
		name, _ := domain.GetName()
		memory, _ := domain.GetMaxMemory()
		id, _ := domain.GetID()
		state, _, _ := domain.GetState()

		var vmState VMState
		if state == libvirt.DOMAIN_RUNNING {
			vmState = Running
		} else {
			vmState = Stopped
		}

		vms = append(vms, &VM{ID: id, Name: name, Memory: memory, State: vmState})
	}
	return vms
}
<|endoftext|>