<commit_before>package teams

import (
	"testing"
	"time"

	"github.com/keybase/client/go/kbtest"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/stretchr/testify/require"

	"golang.org/x/net/context"
)

func TestAccessRequestAccept(t *testing.T) {
	tc, owner, u1, _, teamName := memberSetupMultiple(t)
	defer tc.Cleanup()

	// owner is logged in and created teamName
	err := tc.Logout()
	require.NoError(t, err)

	// u1 requests access to the team
	err = u1.Login(tc.G)
	require.NoError(t, err)
	_, err = RequestAccess(context.Background(), tc.G, teamName)
	require.NoError(t, err)

	myReqs, err := ListMyAccessRequests(context.Background(), tc.G, &teamName)
	require.NoError(t, err)
	require.Equal(t, 1, len(myReqs))
	require.Equal(t, teamName, myReqs[0].String())

	// teamName is optional, if not given, all pending requests will be returned.
	myReqs, err = ListMyAccessRequests(context.Background(), tc.G, nil)
	require.NoError(t, err)
	require.Equal(t, 1, len(myReqs))
	require.Equal(t, teamName, myReqs[0].String())

	// owner lists requests, sees u1 request
	err = tc.Logout()
	require.NoError(t, err)
	err = owner.Login(tc.G)
	require.NoError(t, err)

	reqs, err := ListRequests(context.Background(), tc.G, nil)
	require.NoError(t, err)
	require.Equal(t, 1, len(reqs))
	require.Equal(t, teamName, reqs[0].Name)
	require.Equal(t, u1.Username, reqs[0].Username)
	require.True(t, reqs[0].Ctime.Time().After(time.Now().Add(-1*time.Minute)))
	require.Equal(t, "", reqs[0].FullName.String()) // no fullname in this case

	// owner add u1 to team
	_, err = AddMember(context.Background(), tc.G, teamName, u1.Username, keybase1.TeamRole_WRITER, nil)
	require.NoError(t, err)

	// owner lists requests, sees no requests
	assertNoRequests(tc)

	// u1 requests access to the team again
	err = tc.Logout()
	require.NoError(t, err)
	err = u1.Login(tc.G)
	require.NoError(t, err)

	_, err = RequestAccess(context.Background(), tc.G, teamName)
	require.Error(t, err)
	aerr, ok := err.(libkb.AppStatusError)
	if !ok {
		t.Fatalf("error %s (%T), expected libkb.AppStatusError", err, err)
	}
	if aerr.Code != libkb.SCTeamMemberExists {
		t.Errorf("status code: %d, expected %d", aerr.Code, libkb.SCTeamMemberExists)
	}
	err = tc.Logout()
	require.NoError(t, err)

	// owner lists requests, sees no requests
	err = owner.Login(tc.G)
	require.NoError(t, err)
	assertNoRequests(tc)
}

func TestAccessRequestIgnore(t *testing.T) {
	t.Skip()

	tc, owner, u1, _, teamName := memberSetupMultiple(t)
	defer tc.Cleanup()

	// owner is logged in and created teamName
	err := tc.Logout()
	require.NoError(t, err)

	// u1 requests access to the team
	if err := u1.Login(tc.G); err != nil {
		t.Fatal(err)
	}
	if _, err := RequestAccess(context.Background(), tc.G, teamName); err != nil {
		t.Fatal(err)
	}
	fullName, err := libkb.RandString("test", 5)
	if err != nil {
		t.Fatal(err)
	}
	if err := kbtest.EditProfile(u1.User.MetaContext(context.Background()), keybase1.ProfileEditArg{
		FullName: fullName,
	}); err != nil {
		t.Fatal(err)
	}

	// owner lists requests, sees u1 request
	err = tc.Logout()
	require.NoError(t, err)
	if err := owner.Login(tc.G); err != nil {
		t.Fatal(err)
	}
	reqs, err := ListRequests(context.Background(), tc.G, nil)
	if err != nil {
		t.Fatal(err)
	}
	if len(reqs) != 1 {
		t.Fatalf("num requests: %d, expected 1", len(reqs))
	}
	if reqs[0].Name != teamName {
		t.Errorf("request team name: %q, expected %q", reqs[0].Name, teamName)
	}
	if reqs[0].Username != u1.Username {
		t.Errorf("request username: %q, expected %q", reqs[0].Username, u1.Username)
	}
	if !reqs[0].Ctime.Time().After(time.Now().Add(-1 * time.Minute)) {
		t.Errorf("request ctime %q, expected during last minute", reqs[0].Ctime)
	}
	if reqs[0].FullName.String() != fullName {
		t.Errorf("request full name %q, expected %q", reqs[0].FullName, fullName)
	}

	// owner ignores u1 request
	if err := IgnoreRequest(context.Background(), tc.G, reqs[0].Name, reqs[0].Username); err != nil {
		t.Fatal(err)
	}

	// owner lists requests, sees no requests
	assertNoRequests(tc)

	// u1 requests access to the team again
	err = tc.Logout()
	require.NoError(t, err)
	if err := u1.Login(tc.G); err != nil {
		t.Fatal(err)
	}
	_, err = RequestAccess(context.Background(), tc.G, teamName)
	if err == nil {
		t.Fatal("second RequestAccess success, expected error")
	}
	aerr, ok := err.(libkb.AppStatusError)
	if !ok {
		t.Fatalf("error %s (%T), expected libkb.AppStatusError", err, err)
	}
	if aerr.Code != libkb.SCTeamTarDuplicate {
		t.Errorf("status code: %d, expected %d", aerr.Code, libkb.SCTeamTarDuplicate)
	}
	err = tc.Logout()
	require.NoError(t, err)

	// owner lists requests, sees no requests
	if err := owner.Login(tc.G); err != nil {
		t.Fatal(err)
	}
	assertNoRequests(tc)
}

func assertNoRequests(tc libkb.TestContext) {
	reqs, err := ListRequests(context.Background(), tc.G, nil)
	if err != nil {
		tc.T.Fatal(err)
	}
	if len(reqs) != 0 {
		tc.T.Fatalf("num requests: %d, expected 0", len(reqs))
	}
}
<commit_msg>TAR flake fix (#22508)<commit_after>package teams

import (
	"testing"
	"time"

	"github.com/keybase/client/go/kbtest"

	"github.com/keybase/client/go/libkb"
	"github.com/keybase/client/go/protocol/keybase1"
	"github.com/stretchr/testify/require"

	"golang.org/x/net/context"
)

func TestAccessRequestAccept(t *testing.T) {
	tc, owner, u1, _, teamName := memberSetupMultiple(t)
	defer tc.Cleanup()

	// owner is logged in and created teamName
	err := tc.Logout()
	require.NoError(t, err)

	// u1 requests access to the team
	err = u1.Login(tc.G)
	require.NoError(t, err)
	_, err = RequestAccess(context.Background(), tc.G, teamName)
	require.NoError(t, err)

	myReqs, err := ListMyAccessRequests(context.Background(), tc.G, &teamName)
	require.NoError(t, err)
	require.Equal(t, 1, len(myReqs))
	require.Equal(t, teamName, myReqs[0].String())

	// teamName is optional, if not given, all pending requests will be returned.
	myReqs, err = ListMyAccessRequests(context.Background(), tc.G, nil)
	require.NoError(t, err)
	require.Equal(t, 1, len(myReqs))
	require.Equal(t, teamName, myReqs[0].String())

	// owner lists requests, sees u1 request
	err = tc.Logout()
	require.NoError(t, err)
	err = owner.Login(tc.G)
	require.NoError(t, err)

	reqs, err := ListRequests(context.Background(), tc.G, nil)
	require.NoError(t, err)
	require.Equal(t, 1, len(reqs))
	require.Equal(t, teamName, reqs[0].Name)
	require.Equal(t, u1.Username, reqs[0].Username)
	require.True(t, reqs[0].Ctime.Time().After(time.Now().Add(-1*time.Minute)))
	require.Equal(t, "", reqs[0].FullName.String()) // no fullname in this case

	// owner add u1 to team
	_, err = AddMember(context.Background(), tc.G, teamName, u1.Username, keybase1.TeamRole_WRITER, nil)
	require.NoError(t, err)

	// owner lists requests, sees no requests
	assertNoRequests(tc)

	// u1 requests access to the team again
	err = tc.Logout()
	require.NoError(t, err)
	err = u1.Login(tc.G)
	require.NoError(t, err)

	_, err = RequestAccess(context.Background(), tc.G, teamName)
	require.Error(t, err)
	aerr, ok := err.(libkb.AppStatusError)
	if !ok {
		t.Fatalf("error %s (%T), expected libkb.AppStatusError", err, err)
	}
	if aerr.Code != libkb.SCTeamMemberExists {
		t.Errorf("status code: %d, expected %d", aerr.Code, libkb.SCTeamMemberExists)
	}
	err = tc.Logout()
	require.NoError(t, err)

	// owner lists requests, sees no requests
	err = owner.Login(tc.G)
	require.NoError(t, err)
	assertNoRequests(tc)
}

func TestAccessRequestIgnore(t *testing.T) {
	tc, owner, u1, _, teamName := memberSetupMultiple(t)
	defer tc.Cleanup()

	// owner is logged in and created teamName
	err := tc.Logout()
	require.NoError(t, err)

	// u1 requests access to the team
	if err := u1.Login(tc.G); err != nil {
		t.Fatal(err)
	}
	if _, err := RequestAccess(context.Background(), tc.G, teamName); err != nil {
		t.Fatal(err)
	}
	fullName, err := libkb.RandString("test", 5)
	if err != nil {
		t.Fatal(err)
	}
	if err := kbtest.EditProfile(u1.User.MetaContext(context.Background()), keybase1.ProfileEditArg{
		FullName: fullName,
	}); err != nil {
		t.Fatal(err)
	}
	require.NoError(t, tc.G.UIDMapper.ClearUIDFullName(context.Background(), tc.G, u1.User.GetUID()))

	// owner lists requests, sees u1 request
	err = tc.Logout()
	require.NoError(t, err)
	if err := owner.Login(tc.G); err != nil {
		t.Fatal(err)
	}
	reqs, err := ListRequests(context.Background(), tc.G, nil)
	if err != nil {
		t.Fatal(err)
	}
	if len(reqs) != 1 {
		t.Fatalf("num requests: %d, expected 1", len(reqs))
	}
	if reqs[0].Name != teamName {
		t.Errorf("request team name: %q, expected %q", reqs[0].Name, teamName)
	}
	if reqs[0].Username != u1.Username {
		t.Errorf("request username: %q, expected %q", reqs[0].Username, u1.Username)
	}
	if !reqs[0].Ctime.Time().After(time.Now().Add(-1 * time.Minute)) {
		t.Errorf("request ctime %q, expected during last minute", reqs[0].Ctime)
	}
	if reqs[0].FullName.String() != fullName {
		t.Errorf("request full name %q, expected %q", reqs[0].FullName, fullName)
	}

	// owner ignores u1 request
	if err := IgnoreRequest(context.Background(), tc.G, reqs[0].Name, reqs[0].Username); err != nil {
		t.Fatal(err)
	}

	// owner lists requests, sees no requests
	assertNoRequests(tc)

	// u1 requests access to the team again
	err = tc.Logout()
	require.NoError(t, err)
	if err := u1.Login(tc.G); err != nil {
		t.Fatal(err)
	}
	_, err = RequestAccess(context.Background(), tc.G, teamName)
	if err == nil {
		t.Fatal("second RequestAccess success, expected error")
	}
	aerr, ok := err.(libkb.AppStatusError)
	if !ok {
		t.Fatalf("error %s (%T), expected libkb.AppStatusError", err, err)
	}
	if aerr.Code != libkb.SCTeamTarDuplicate {
		t.Errorf("status code: %d, expected %d", aerr.Code, libkb.SCTeamTarDuplicate)
	}
	err = tc.Logout()
	require.NoError(t, err)

	// owner lists requests, sees no requests
	if err := owner.Login(tc.G); err != nil {
		t.Fatal(err)
	}
	assertNoRequests(tc)
}

func assertNoRequests(tc libkb.TestContext) {
	reqs, err := ListRequests(context.Background(), tc.G, nil)
	if err != nil {
		tc.T.Fatal(err)
	}
	if len(reqs) != 0 {
		tc.T.Fatalf("num requests: %d, expected 0", len(reqs))
	}
}
<|endoftext|>


<commit_before>package notification

import (
	"errors"
	"github.com/jinzhu/gorm"
	"github.com/koding/bongo"
	"github.com/koding/logging"
	"math"
	"net/http"
	"net/url"
	"socialapi/config"
	"socialapi/models"
	"socialapi/workers/api/modules/helpers"
	"socialapi/workers/cache"
	"socialapi/workers/helper"
	"strconv"
)

var (
	NOTIFICATION_LIMIT = 8
	ACTOR_LIMIT        = 3
	cacheEnabled       = false
	log                logging.Logger
)

func init() {
	log = helper.CreateLogger("NotificationAPI", false)
	helpers.Log = log
}

func List(u *url.URL, h http.Header, _ interface{}) (int, http.Header, interface{}, error) {
	q := helpers.GetQuery(u)
	if err := validateNotificationRequest(q); err != nil {
		return helpers.NewBadRequestResponse(err)
	}

	conf := config.Get()
	cacheEnabled = conf.Cache.Notification

	urlQuery := u.Query()
	cache, err := strconv.ParseBool(urlQuery.Get("cache"))
	if err == nil {
		cacheEnabled = cache
	}

	list, err := fetchNotifications(q)
	if err != nil {
		if err == gorm.RecordNotFound {
			return helpers.NewNotFoundResponse()
		}

		return helpers.NewBadRequestResponse(err)
	}

	return helpers.NewOKResponse(list)
}

func Glance(u *url.URL, h http.Header, req *models.Notification) (int, http.Header, interface{}, error) {
	q := models.NewQuery()
	q.AccountId = req.AccountId
	if err := validateNotificationRequest(q); err != nil {
		return helpers.NewBadRequestResponse(err)
	}

	if err := req.Glance(); err != nil {
		if err == gorm.RecordNotFound {
			return helpers.NewNotFoundResponse()
		}

		return helpers.NewBadRequestResponse(err)
	}

	go func() {
		if cacheEnabled {
			cacheInstance := cache.NewNotificationCache()
			cacheInstance.Glance(req)
		}
	}()

	req.Glanced = true

	return helpers.NewDefaultOKResponse()
}

func Follow(u *url.URL, h http.Header, req *models.Activity) (int, http.Header, interface{}, error) {

	n := models.NewNotification()
	if err := n.Follow(req); err != nil {
		return helpers.NewBadRequestResponse(err)
	}

	return helpers.NewDefaultOKResponse()
}

type GroupRequest struct {
	Name         string  `json:"name"`
	TypeConstant string  `json:"typeConstant"`
	ActorId      int64   `json:"actorId"`
	Admins       []int64 `json:"admins"`
}

func InteractGroup(u *url.URL, h http.Header, req *GroupRequest) (int, http.Header, interface{}, error) {

	// first fetch channel id as target id
	c := models.NewChannel()
	selector := map[string]interface{}{
		"type_constant": models.Channel_TYPE_GROUP,
		"group_name":    req.Name,
		"name":          req.Name,
	}

	if err := c.One(bongo.NewQS(selector)); err != nil {
		return helpers.NewBadRequestResponse(err)
	}

	a := models.NewActivity()
	a.TargetId = c.Id
	a.ActorId = req.ActorId
	a.TypeConstant = req.TypeConstant

	var err error

	n := models.NewNotification()
	switch req.TypeConstant {
	case models.NotificationContent_TYPE_JOIN:
		err = n.JoinGroup(a, req.Admins)
	case models.NotificationContent_TYPE_LEAVE:
		err = n.LeaveGroup(a, req.Admins)
	default:
		err = errors.New("group interaction type not found")
	}
	if err != nil {
		return helpers.NewBadRequestResponse(err)
	}

	return helpers.NewDefaultOKResponse()
}

func fetchNotifications(q *models.Query) (*models.NotificationResponse, error) {
	var list *models.NotificationResponse
	var err error

	// first check redis
	var cacheInstance *cache.NotificationCache
	if cacheEnabled {
		cacheInstance = cache.NewNotificationCache()
		cacheInstance.ActorLimit = ACTOR_LIMIT
		list, err = cacheInstance.FetchNotifications(q.AccountId)
		if err != nil {
			log.Error("Cache error: %s", err)
		}
		if err == nil && len(list.Notifications) > 0 {
			return list, nil
		}
	}

	n := models.NewNotification()
	list, err = n.List(q)
	if err != nil {
		return nil, err
	}

	go func() {
		if cacheEnabled {
			if err := cacheInstance.UpdateCachedNotifications(q.AccountId, list); err != nil {
				log.Error("Cache cannot be updated: %s", err)
			}
		}
	}()

	return list, nil
}

func validateNotificationRequest(q *models.Query) error {
	a := models.NewAccount()
	a.Id = q.AccountId
	if q.AccountId == 0 {
		return errors.New("Account id cannot be empty")
	}

	if err := a.Fetch(); err != nil {
		if err == gorm.RecordNotFound {
			return errors.New("Account is not valid")
		}
		return err
	}

	// reset the limit if it is needed
	q.Limit = int(math.Min(float64(q.Limit), float64(NOTIFICATION_LIMIT)))

	return nil
}
<commit_msg>Social: validateSubscriptionRequest is added and account validation lines in validateNotificationRequest is taken into validateAccount method in notification api<commit_after>package notification

import (
	"errors"
	"github.com/jinzhu/gorm"
	"github.com/koding/bongo"
	"github.com/koding/logging"
	"math"
	"net/http"
	"net/url"
	"socialapi/config"
	"socialapi/models"
	"socialapi/workers/api/modules/helpers"
	"socialapi/workers/cache"
	"socialapi/workers/helper"
	"strconv"
)

var (
	NOTIFICATION_LIMIT = 8
	ACTOR_LIMIT        = 3
	cacheEnabled       = false
	log                logging.Logger
)

func init() {
	log = helper.CreateLogger("NotificationAPI", false)
	helpers.Log = log
}

func List(u *url.URL, h http.Header, _ interface{}) (int, http.Header, interface{}, error) {
	q := helpers.GetQuery(u)
	if err := validateNotificationRequest(q); err != nil {
		return helpers.NewBadRequestResponse(err)
	}

	conf := config.Get()
	cacheEnabled = conf.Cache.Notification

	urlQuery := u.Query()
	cache, err := strconv.ParseBool(urlQuery.Get("cache"))
	if err == nil {
		cacheEnabled = cache
	}

	list, err := fetchNotifications(q)
	if err != nil {
		if err == gorm.RecordNotFound {
			return helpers.NewNotFoundResponse()
		}

		return helpers.NewBadRequestResponse(err)
	}

	return helpers.NewOKResponse(list)
}

func Glance(u *url.URL, h http.Header, req *models.Notification) (int, http.Header, interface{}, error) {
	q := models.NewQuery()
	q.AccountId = req.AccountId
	if err := validateNotificationRequest(q); err != nil {
		return helpers.NewBadRequestResponse(err)
	}

	if err := req.Glance(); err != nil {
		if err == gorm.RecordNotFound {
			return helpers.NewNotFoundResponse()
		}

		return helpers.NewBadRequestResponse(err)
	}

	go func() {
		if cacheEnabled {
			cacheInstance := cache.NewNotificationCache()
			cacheInstance.Glance(req)
		}
	}()

	req.Glanced = true

	return helpers.NewDefaultOKResponse()
}

func Follow(u *url.URL, h http.Header, req *models.Activity) (int, http.Header, interface{}, error) {

	n := models.NewNotification()
	if err := n.Follow(req); err != nil {
		return helpers.NewBadRequestResponse(err)
	}

	return helpers.NewDefaultOKResponse()
}

type GroupRequest struct {
	Name         string  `json:"name"`
	TypeConstant string  `json:"typeConstant"`
	ActorId      int64   `json:"actorId"`
	Admins       []int64 `json:"admins"`
}

func InteractGroup(u *url.URL, h http.Header, req *GroupRequest) (int, http.Header, interface{}, error) {

	// first fetch channel id as target id
	c := models.NewChannel()
	selector := map[string]interface{}{
		"type_constant": models.Channel_TYPE_GROUP,
		"group_name":    req.Name,
		"name":          req.Name,
	}

	if err := c.One(bongo.NewQS(selector)); err != nil {
		return helpers.NewBadRequestResponse(err)
	}

	a := models.NewActivity()
	a.TargetId = c.Id
	a.ActorId = req.ActorId
	a.TypeConstant = req.TypeConstant

	var err error

	n := models.NewNotification()
	switch req.TypeConstant {
	case models.NotificationContent_TYPE_JOIN:
		err = n.JoinGroup(a, req.Admins)
	case models.NotificationContent_TYPE_LEAVE:
		err = n.LeaveGroup(a, req.Admins)
	default:
		err = errors.New("group interaction type not found")
	}
	if err != nil {
		return helpers.NewBadRequestResponse(err)
	}

	return helpers.NewDefaultOKResponse()
}

func fetchNotifications(q *models.Query) (*models.NotificationResponse, error) {
	var list *models.NotificationResponse
	var err error

	// first check redis
	var cacheInstance *cache.NotificationCache
	if cacheEnabled {
		cacheInstance = cache.NewNotificationCache()
		cacheInstance.ActorLimit = ACTOR_LIMIT
		list, err = cacheInstance.FetchNotifications(q.AccountId)
		if err != nil {
			log.Error("Cache error: %s", err)
		}
		if err == nil && len(list.Notifications) > 0 {
			return list, nil
		}
	}

	n := models.NewNotification()
	list, err = n.List(q)
	if err != nil {
		return nil, err
	}

	go func() {
		if cacheEnabled {
			if err := cacheInstance.UpdateCachedNotifications(q.AccountId, list); err != nil {
				log.Error("Cache cannot be updated: %s", err)
			}
		}
	}()

	return list, nil
}

func validateNotificationRequest(q *models.Query) error {
	if err := validateAccount(q.AccountId); err != nil {
		return err
	}
	// update the limit if it is needed
	q.Limit = int(math.Min(float64(q.Limit), float64(NOTIFICATION_LIMIT)))

	return nil
}

func validateSubscriptionRequest(a *models.Activity) error {
	if err := validateAccount(a.ActorId); err != nil {
		return err
	}

	return validateMessage(a.TargetId)
}

func validateAccount(accountId int64) error {
	a := models.NewAccount()
	if accountId == 0 {
		return errors.New("Account id cannot be empty")
	}

	if err := a.ById(accountId); err != nil {
		if err == gorm.RecordNotFound {
			return errors.New("Account is not valid")
		}
		return err
	}

	return nil
}

func validateMessage(messageId int64) error {
	cm := models.NewChannelMessage()
	if err := cm.ById(messageId); err != nil {
		if err == gorm.RecordNotFound {
			return errors.New("Channel message does not exist")
		}
		return err
	}

	return nil
}
<|endoftext|>


<commit_before>package arm

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

	"github.com/Azure/go-autorest/autorest"
	"github.com/Azure/go-autorest/autorest/azure"
	"github.com/Azure/packer-azure/packer/builder/azure/common/logutil"
	"io"
)

func chop(data []byte, maxlen int64) string {
	s := string(data)
	if int64(len(s)) > maxlen {
		s = s[:maxlen] + "..."
	}
	return s
}

func handleBody(body io.ReadCloser, maxlen int64) (io.ReadCloser, string) {
	if body == nil {
		return nil, ""
	}

	b, err := ioutil.ReadAll(body)
	if err != nil {
		return nil, ""
	}

	return ioutil.NopCloser(bytes.NewReader(b)), chop(b, maxlen)
}

func withInspection(maxlen int64) autorest.PrepareDecorator {
	return func(p autorest.Preparer) autorest.Preparer {
		return autorest.PreparerFunc(func(r *http.Request) (*http.Request, error) {
			body, bodyString := handleBody(r.Body, maxlen)
			r.Body = body

			log.Print("Azure request", logutil.Fields{
				"method":  r.Method,
				"request": r.URL.String(),
				"body":    bodyString,
			})
			return p.Prepare(r)
		})
	}
}

func byInspecting(maxlen int64) autorest.RespondDecorator {
	return func(r autorest.Responder) autorest.Responder {
		return autorest.ResponderFunc(func(resp *http.Response) error {
			body, bodyString := handleBody(resp.Body, maxlen)
			resp.Body = body

			log.Print("Azure response", logutil.Fields{
				"status":          resp.Status,
				"method":          resp.Request.Method,
				"request":         resp.Request.URL.String(),
				"x-ms-request-id": azure.ExtractRequestID(resp),
				"body":            bodyString,
			})
			return r.Respond(resp)
		})
	}
}
<commit_msg>Ensure the body is closed.<commit_after>package arm

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

	"github.com/Azure/go-autorest/autorest"
	"github.com/Azure/go-autorest/autorest/azure"
	"github.com/Azure/packer-azure/packer/builder/azure/common/logutil"
	"io"
)

func chop(data []byte, maxlen int64) string {
	s := string(data)
	if int64(len(s)) > maxlen {
		s = s[:maxlen] + "..."
	}
	return s
}

func handleBody(body io.ReadCloser, maxlen int64) (io.ReadCloser, string) {
	defer body.Close()

	if body == nil {
		return nil, ""
	}

	b, err := ioutil.ReadAll(body)
	if err != nil {
		return nil, ""
	}

	return ioutil.NopCloser(bytes.NewReader(b)), chop(b, maxlen)
}

func withInspection(maxlen int64) autorest.PrepareDecorator {
	return func(p autorest.Preparer) autorest.Preparer {
		return autorest.PreparerFunc(func(r *http.Request) (*http.Request, error) {
			body, bodyString := handleBody(r.Body, maxlen)
			r.Body = body

			log.Print("Azure request", logutil.Fields{
				"method":  r.Method,
				"request": r.URL.String(),
				"body":    bodyString,
			})
			return p.Prepare(r)
		})
	}
}

func byInspecting(maxlen int64) autorest.RespondDecorator {
	return func(r autorest.Responder) autorest.Responder {
		return autorest.ResponderFunc(func(resp *http.Response) error {
			body, bodyString := handleBody(resp.Body, maxlen)
			resp.Body = body

			log.Print("Azure response", logutil.Fields{
				"status":          resp.Status,
				"method":          resp.Request.Method,
				"request":         resp.Request.URL.String(),
				"x-ms-request-id": azure.ExtractRequestID(resp),
				"body":            bodyString,
			})
			return r.Respond(resp)
		})
	}
}
<|endoftext|>


<commit_before>package cmd

import (
	"text/template"

	"github.com/spf13/cobra"
)

var versionCmd = &cobra.Command{
	Use:   "version",
	Short: "Print the version number of File Browser",
	Long:  `All software has versions. This is File Browser's`,
	Run: func(cmd *cobra.Command, args []string) {
		// https://github.com/spf13/cobra/issues/724
		t := template.New("version")
		template.Must(t.Parse(rootCmd.VersionTemplate()))
		err := t.Execute(rootCmd.OutOrStdout(), rootCmd)
		if err != nil {
			rootCmd.Println(err)
		}
	},
}

func init() {
	rootCmd.AddCommand(versionCmd)
}
<commit_msg>chore: add version to subcmds<commit_after>package cmd

import (
	"text/template"

	"github.com/spf13/cobra"
)

var versionCmd = &cobra.Command{
	Use:   "version",
	Short: "Print the version number of File Browser",
	Long:  `All software has versions. This is File Browser's`,
	Run: func(cmd *cobra.Command, args []string) {
		// https://github.com/spf13/cobra/issues/724
		t := template.New("version")
		template.Must(t.Parse(rootCmd.VersionTemplate()))
		err := t.Execute(rootCmd.OutOrStdout(), rootCmd)
		if err != nil {
			rootCmd.Println(err)
		}
	},
}

func init() {
	rootCmd.AddCommand(versionCmd)
	serveCmd.AddCommand(versionCmd)
	dbCmd.AddCommand(versionCmd)
}
<|endoftext|>


<commit_before>package command

import (
	"io"
	"net/http"
	"os"
	"runtime"

	"github.com/docker/cli/cli"
	cliconfig "github.com/docker/cli/cli/config"
	"github.com/docker/cli/cli/config/configfile"
	cliflags "github.com/docker/cli/cli/flags"
	dopts "github.com/docker/cli/opts"
	"github.com/docker/docker/api"
	"github.com/docker/docker/client"
	"github.com/docker/go-connections/sockets"
	"github.com/docker/go-connections/tlsconfig"
	"github.com/docker/notary/passphrase"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"golang.org/x/net/context"
)

// Streams is an interface which exposes the standard input and output streams
type Streams interface {
	In() *InStream
	Out() *OutStream
	Err() io.Writer
}

// Cli represents the docker command line client.
type Cli interface {
	Client() client.APIClient
	Out() *OutStream
	Err() io.Writer
	In() *InStream
	SetIn(in *InStream)
	ConfigFile() *configfile.ConfigFile
	ServerInfo() ServerInfo
}

// DockerCli is an instance the docker command line client.
// Instances of the client can be returned from NewDockerCli.
type DockerCli struct {
	configFile     *configfile.ConfigFile
	in             *InStream
	out            *OutStream
	err            io.Writer
	client         client.APIClient
	defaultVersion string
	server         ServerInfo
}

// DefaultVersion returns api.defaultVersion or DOCKER_API_VERSION if specified.
func (cli *DockerCli) DefaultVersion() string {
	return cli.defaultVersion
}

// Client returns the APIClient
func (cli *DockerCli) Client() client.APIClient {
	return cli.client
}

// Out returns the writer used for stdout
func (cli *DockerCli) Out() *OutStream {
	return cli.out
}

// Err returns the writer used for stderr
func (cli *DockerCli) Err() io.Writer {
	return cli.err
}

// SetIn sets the reader used for stdin
func (cli *DockerCli) SetIn(in *InStream) {
	cli.in = in
}

// In returns the reader used for stdin
func (cli *DockerCli) In() *InStream {
	return cli.in
}

// ShowHelp shows the command help.
func ShowHelp(err io.Writer) func(*cobra.Command, []string) error {
	return func(cmd *cobra.Command, args []string) error {
		cmd.SetOutput(err)
		cmd.HelpFunc()(cmd, args)
		return nil
	}
}

// ConfigFile returns the ConfigFile
func (cli *DockerCli) ConfigFile() *configfile.ConfigFile {
	return cli.configFile
}

// ServerInfo returns the server version details for the host this client is
// connected to
func (cli *DockerCli) ServerInfo() ServerInfo {
	return cli.server
}

// Initialize the dockerCli runs initialization that must happen after command
// line flags are parsed.
func (cli *DockerCli) Initialize(opts *cliflags.ClientOptions) error {
	cli.configFile = cliconfig.LoadDefaultConfigFile(cli.err)

	var err error
	cli.client, err = NewAPIClientFromFlags(opts.Common, cli.configFile)
	if tlsconfig.IsErrEncryptedKey(err) {
		var (
			passwd string
			giveup bool
		)
		passRetriever := passphrase.PromptRetrieverWithInOut(cli.In(), cli.Out(), nil)

		for attempts := 0; tlsconfig.IsErrEncryptedKey(err); attempts++ {
			// some code and comments borrowed from notary/trustmanager/keystore.go
			passwd, giveup, err = passRetriever("private", "encrypted TLS private", false, attempts)
			// Check if the passphrase retriever got an error or if it is telling us to give up
			if giveup || err != nil {
				return errors.Wrap(err, "private key is encrypted, but could not get passphrase")
			}

			opts.Common.TLSOptions.Passphrase = passwd
			cli.client, err = NewAPIClientFromFlags(opts.Common, cli.configFile)
		}
	}

	if err != nil {
		return err
	}

	cli.defaultVersion = cli.client.ClientVersion()

	if ping, err := cli.client.Ping(context.Background()); err == nil {
		cli.server = ServerInfo{
			HasExperimental: ping.Experimental,
			OSType:          ping.OSType,
		}

		cli.client.NegotiateAPIVersionPing(ping)
	} else {
		// Default to true if we fail to connect to daemon
		cli.server = ServerInfo{HasExperimental: true}
	}

	return nil
}

// ServerInfo stores details about the supported features and platform of the
// server
type ServerInfo struct {
	HasExperimental bool
	OSType          string
}

// NewDockerCli returns a DockerCli instance with IO output and error streams set by in, out and err.
func NewDockerCli(in io.ReadCloser, out, err io.Writer) *DockerCli {
	return &DockerCli{in: NewInStream(in), out: NewOutStream(out), err: err}
}

// NewAPIClientFromFlags creates a new APIClient from command line flags
func NewAPIClientFromFlags(opts *cliflags.CommonOptions, configFile *configfile.ConfigFile) (client.APIClient, error) {
	host, err := getServerHost(opts.Hosts, opts.TLSOptions)
	if err != nil {
		return &client.Client{}, err
	}

	customHeaders := configFile.HTTPHeaders
	if customHeaders == nil {
		customHeaders = map[string]string{}
	}
	customHeaders["User-Agent"] = UserAgent()

	verStr := api.DefaultVersion
	if tmpStr := os.Getenv("DOCKER_API_VERSION"); tmpStr != "" {
		verStr = tmpStr
	}

	httpClient, err := newHTTPClient(host, opts.TLSOptions)
	if err != nil {
		return &client.Client{}, err
	}

	return client.NewClient(host, verStr, httpClient, customHeaders)
}

func getServerHost(hosts []string, tlsOptions *tlsconfig.Options) (host string, err error) {
	switch len(hosts) {
	case 0:
		host = os.Getenv("DOCKER_HOST")
	case 1:
		host = hosts[0]
	default:
		return "", errors.New("Please specify only one -H")
	}

	host, err = dopts.ParseHost(tlsOptions != nil, host)
	return
}

func newHTTPClient(host string, tlsOptions *tlsconfig.Options) (*http.Client, error) {
	if tlsOptions == nil {
		// let the api client configure the default transport.
		return nil, nil
	}
	opts := *tlsOptions
	opts.ExclusiveRootPools = true
	config, err := tlsconfig.Client(opts)
	if err != nil {
		return nil, err
	}
	tr := &http.Transport{
		TLSClientConfig: config,
	}
	proto, addr, _, err := client.ParseHost(host)
	if err != nil {
		return nil, err
	}

	sockets.ConfigureTransport(tr, proto, addr)

	return &http.Client{
		Transport:     tr,
		CheckRedirect: client.CheckRedirect,
	}, nil
}

// UserAgent returns the user agent string used for making API requests
func UserAgent() string {
	return "Docker-Client/" + cli.Version + " (" + runtime.GOOS + ")"
}
<commit_msg>Enable TCP Keep-Alive in Docker client<commit_after>package command

import (
	"io"
	"net"
	"net/http"
	"os"
	"runtime"
	"time"

	"github.com/docker/cli/cli"
	cliconfig "github.com/docker/cli/cli/config"
	"github.com/docker/cli/cli/config/configfile"
	cliflags "github.com/docker/cli/cli/flags"
	dopts "github.com/docker/cli/opts"
	"github.com/docker/docker/api"
	"github.com/docker/docker/client"
	"github.com/docker/go-connections/sockets"
	"github.com/docker/go-connections/tlsconfig"
	"github.com/docker/notary/passphrase"
	"github.com/pkg/errors"
	"github.com/spf13/cobra"
	"golang.org/x/net/context"
)

// Streams is an interface which exposes the standard input and output streams
type Streams interface {
	In() *InStream
	Out() *OutStream
	Err() io.Writer
}

// Cli represents the docker command line client.
type Cli interface {
	Client() client.APIClient
	Out() *OutStream
	Err() io.Writer
	In() *InStream
	SetIn(in *InStream)
	ConfigFile() *configfile.ConfigFile
	ServerInfo() ServerInfo
}

// DockerCli is an instance the docker command line client.
// Instances of the client can be returned from NewDockerCli.
type DockerCli struct {
	configFile     *configfile.ConfigFile
	in             *InStream
	out            *OutStream
	err            io.Writer
	client         client.APIClient
	defaultVersion string
	server         ServerInfo
}

// DefaultVersion returns api.defaultVersion or DOCKER_API_VERSION if specified.
func (cli *DockerCli) DefaultVersion() string {
	return cli.defaultVersion
}

// Client returns the APIClient
func (cli *DockerCli) Client() client.APIClient {
	return cli.client
}

// Out returns the writer used for stdout
func (cli *DockerCli) Out() *OutStream {
	return cli.out
}

// Err returns the writer used for stderr
func (cli *DockerCli) Err() io.Writer {
	return cli.err
}

// SetIn sets the reader used for stdin
func (cli *DockerCli) SetIn(in *InStream) {
	cli.in = in
}

// In returns the reader used for stdin
func (cli *DockerCli) In() *InStream {
	return cli.in
}

// ShowHelp shows the command help.
func ShowHelp(err io.Writer) func(*cobra.Command, []string) error {
	return func(cmd *cobra.Command, args []string) error {
		cmd.SetOutput(err)
		cmd.HelpFunc()(cmd, args)
		return nil
	}
}

// ConfigFile returns the ConfigFile
func (cli *DockerCli) ConfigFile() *configfile.ConfigFile {
	return cli.configFile
}

// ServerInfo returns the server version details for the host this client is
// connected to
func (cli *DockerCli) ServerInfo() ServerInfo {
	return cli.server
}

// Initialize the dockerCli runs initialization that must happen after command
// line flags are parsed.
func (cli *DockerCli) Initialize(opts *cliflags.ClientOptions) error {
	cli.configFile = cliconfig.LoadDefaultConfigFile(cli.err)

	var err error
	cli.client, err = NewAPIClientFromFlags(opts.Common, cli.configFile)
	if tlsconfig.IsErrEncryptedKey(err) {
		var (
			passwd string
			giveup bool
		)
		passRetriever := passphrase.PromptRetrieverWithInOut(cli.In(), cli.Out(), nil)

		for attempts := 0; tlsconfig.IsErrEncryptedKey(err); attempts++ {
			// some code and comments borrowed from notary/trustmanager/keystore.go
			passwd, giveup, err = passRetriever("private", "encrypted TLS private", false, attempts)
			// Check if the passphrase retriever got an error or if it is telling us to give up
			if giveup || err != nil {
				return errors.Wrap(err, "private key is encrypted, but could not get passphrase")
			}

			opts.Common.TLSOptions.Passphrase = passwd
			cli.client, err = NewAPIClientFromFlags(opts.Common, cli.configFile)
		}
	}

	if err != nil {
		return err
	}

	cli.defaultVersion = cli.client.ClientVersion()

	if ping, err := cli.client.Ping(context.Background()); err == nil {
		cli.server = ServerInfo{
			HasExperimental: ping.Experimental,
			OSType:          ping.OSType,
		}

		cli.client.NegotiateAPIVersionPing(ping)
	} else {
		// Default to true if we fail to connect to daemon
		cli.server = ServerInfo{HasExperimental: true}
	}

	return nil
}

// ServerInfo stores details about the supported features and platform of the
// server
type ServerInfo struct {
	HasExperimental bool
	OSType          string
}

// NewDockerCli returns a DockerCli instance with IO output and error streams set by in, out and err.
func NewDockerCli(in io.ReadCloser, out, err io.Writer) *DockerCli {
	return &DockerCli{in: NewInStream(in), out: NewOutStream(out), err: err}
}

// NewAPIClientFromFlags creates a new APIClient from command line flags
func NewAPIClientFromFlags(opts *cliflags.CommonOptions, configFile *configfile.ConfigFile) (client.APIClient, error) {
	host, err := getServerHost(opts.Hosts, opts.TLSOptions)
	if err != nil {
		return &client.Client{}, err
	}

	customHeaders := configFile.HTTPHeaders
	if customHeaders == nil {
		customHeaders = map[string]string{}
	}
	customHeaders["User-Agent"] = UserAgent()

	verStr := api.DefaultVersion
	if tmpStr := os.Getenv("DOCKER_API_VERSION"); tmpStr != "" {
		verStr = tmpStr
	}

	httpClient, err := newHTTPClient(host, opts.TLSOptions)
	if err != nil {
		return &client.Client{}, err
	}

	return client.NewClient(host, verStr, httpClient, customHeaders)
}

func getServerHost(hosts []string, tlsOptions *tlsconfig.Options) (host string, err error) {
	switch len(hosts) {
	case 0:
		host = os.Getenv("DOCKER_HOST")
	case 1:
		host = hosts[0]
	default:
		return "", errors.New("Please specify only one -H")
	}

	host, err = dopts.ParseHost(tlsOptions != nil, host)
	return
}

func newHTTPClient(host string, tlsOptions *tlsconfig.Options) (*http.Client, error) {
	if tlsOptions == nil {
		// let the api client configure the default transport.
		return nil, nil
	}
	opts := *tlsOptions
	opts.ExclusiveRootPools = true
	config, err := tlsconfig.Client(opts)
	if err != nil {
		return nil, err
	}
	tr := &http.Transport{
		TLSClientConfig: config,
		DialContext: (&net.Dialer{
			KeepAlive: 30 * time.Second,
			Timeout:   30 * time.Second,
		}).DialContext,
	}
	proto, addr, _, err := client.ParseHost(host)
	if err != nil {
		return nil, err
	}

	sockets.ConfigureTransport(tr, proto, addr)

	return &http.Client{
		Transport:     tr,
		CheckRedirect: client.CheckRedirect,
	}, nil
}

// UserAgent returns the user agent string used for making API requests
func UserAgent() string {
	return "Docker-Client/" + cli.Version + " (" + runtime.GOOS + ")"
}
<|endoftext|>


<commit_before>// Package encoding converts arbitrary objects into byte slices, and vis
// versa. It also contains helper functions for reading and writing length-
// prefixed data. The encoding rules are as follows:
//
// Objects are encoded as binary data, without type information. The receiver
// uses context to determine the type to decode into.
//
// Integers are little-endian, and are always encoded as 8 bytes, i.e. their
// int64 or uint64 equivalent.
//
// Booleans are encoded as one byte, either zero (false) or one (true). No
// other values may be used.
//
// Nil pointers are equivalent to "false," i.e. a single zero byte. Valid
// pointers are represented by a "true" byte (0x01) followed by the encoding
// of the dereferenced value.
//
// Variable-length types, such as strings and slices, are represented by an 8-byte
// length-prefix followed by the encoded value.
//
// Slices and structs are simply the concatenation of their encoded elements.
// Byte slices are not subject to the 8-byte integer rule; they are encoded as
// their literal representation, one byte per byte.
//
// The ordering of struct fields is determined by their type definition. For
// example:
//
//   type foo struct { S string; I int }
//
//   Marshal(foo{"bar", 3}) == append(Marshal("bar"), Marshal(3)...)
//
// Finally, if a type implements the SiaMarshaler interface, its MarshalSia
// method will be used to encode the type. The resulting byte slice will be
// length-prefixed like any other variable-length type. During decoding, the
// type is decoded as a byte slice, and then passed to UnmarshalSia.
package encoding

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"reflect"
)

// A SiaMarshaler can encode itself as a byte slice.
type SiaMarshaler interface {
	MarshalSia() []byte
}

// A SiaUnmarshaler can decode itself from a byte slice. If a decoding error
// occurs, UnmarshalSia should panic.
type SiaUnmarshaler interface {
	UnmarshalSia([]byte)
}

// An Encoder writes objects to an output stream.
type Encoder struct {
	w io.Writer
}

const (
	maxSliceLen = 4 * 1024 * 1024 // 4 MB
)

var (
	errBadPointer = errors.New("cannot decode into invalid pointer")
)

// Encode writes the encoding of v to the stream. For encoding details, see
// the package docstring.
func (e *Encoder) Encode(v interface{}) error {
	return e.encode(reflect.ValueOf(v))
}

// write catches instances where short writes do not return an error.
func (e *Encoder) write(p []byte) error {
	n, err := e.w.Write(p)
	if n != len(p) {
		return io.ErrShortWrite
	}
	return err
}

// Encode writes the encoding of val to the stream. For encoding details, see
// the package docstring.
func (e *Encoder) encode(val reflect.Value) error {
	// check for MarshalSia interface first
	if m, ok := val.Interface().(SiaMarshaler); ok {
		return WritePrefix(e.w, m.MarshalSia())
	}

	switch val.Kind() {
	case reflect.Ptr:
		// write either a 1 or 0
		if err := e.Encode(!val.IsNil()); err != nil {
			return err
		}
		if !val.IsNil() {
			return e.encode(val.Elem())
		}
	case reflect.Bool:
		if val.Bool() {
			return e.write([]byte{1})
		} else {
			return e.write([]byte{0})
		}
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return e.write(EncInt64(val.Int()))
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return e.write(EncUint64(val.Uint()))
	case reflect.String:
		return WritePrefix(e.w, []byte(val.String()))
	case reflect.Slice:
		// slices are variable length, so prepend the length and then fallthrough to array logic
		if err := e.write(EncUint64(uint64(val.Len()))); err != nil {
			return err
		}
		fallthrough
	case reflect.Array:
		// special case for byte arrays
		if val.Type().Elem().Kind() == reflect.Uint8 {
			// convert array to slice so we can use Bytes()
			// can't just use Slice() because array may be unaddressable
			slice := reflect.MakeSlice(reflect.SliceOf(val.Type().Elem()), val.Len(), val.Len())
			reflect.Copy(slice, val)
			return e.write(slice.Bytes())
		}
		// normal slices/arrays are encoded by sequentially encoding their elements
		for i := 0; i < val.Len(); i++ {
			if err := e.encode(val.Index(i)); err != nil {
				return err
			}
		}
	case reflect.Struct:
		for i := 0; i < val.NumField(); i++ {
			if err := e.encode(val.Field(i)); err != nil {
				return err
			}
		}
	default:
		// Marshalling should never fail. If it panics, you're doing something wrong,
		// like trying to encode a map or an unexported struct field.
		panic("could not marshal type " + val.Type().String())
	}
	return nil
}

// NewEncoder returns a new encoder that writes to w.
func NewEncoder(w io.Writer) *Encoder {
	return &Encoder{w}
}

// Marshal returns the encoding of v. For encoding details, see the package
// docstring.
func Marshal(v interface{}) []byte {
	b := new(bytes.Buffer)
	NewEncoder(b).Encode(v) // no error possible when using a bytes.Buffer
	return b.Bytes()
}

// MarshalAll encodes all of its inputs and returns their concatenation.
func MarshalAll(v ...interface{}) []byte {
	b := new(bytes.Buffer)
	enc := NewEncoder(b)
	for i := range v {
		enc.Encode(v[i]) // no error possible when using a bytes.Buffer
	}
	return b.Bytes()
}

// WriteFile writes v to a file. The file will be created if it does not exist.
func WriteFile(filename string, v interface{}) error {
	file, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE, 0660)
	if err != nil {
		return err
	}
	defer file.Close()
	return NewEncoder(file).Encode(v)
}

// A Decoder reads and decodes values from an input stream.
type Decoder struct {
	r io.Reader
}

// Decode reads the next encoded value from its input stream and stores it in
// v, which must be a pointer. The decoding rules are the inverse of those
// specified in the package docstring.
func (d *Decoder) Decode(v interface{}) (err error) {
	// v must be a pointer
	pval := reflect.ValueOf(v)
	if pval.Kind() != reflect.Ptr || pval.IsNil() {
		return errBadPointer
	}

	// catch decoding panics and convert them to errors
	// note that this allows us to skip boundary checks during decoding
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("could not decode type %s: %v", pval.Elem().Type().String(), r)
		}
	}()

	d.decode(pval.Elem())
	return
}

// readN reads n bytes and panics if the read fails.
func (d *Decoder) readN(n int) []byte {
	b := make([]byte, n)
	_, err := io.ReadFull(d.r, b)
	if err != nil {
		panic(err)
	}
	return b
}

// readPrefix reads a length-prefixed byte slice and panics if the read fails.
func (d *Decoder) readPrefix() []byte {
	// TODO: what should maxlen be?
	b, err := ReadPrefix(d.r, 1<<32)
	if err != nil {
		panic(err)
	}
	return b
}

// decode reads the next encoded value from its input stream and stores it in
// val. The decoding rules are the inverse of those specified in the package
// docstring.
func (d *Decoder) decode(val reflect.Value) {
	// check for UnmarshalSia interface first
	if val.CanAddr() {
		if u, ok := val.Addr().Interface().(SiaUnmarshaler); ok {
			u.UnmarshalSia(d.readPrefix())
			return
		}
	}

	switch val.Kind() {
	case reflect.Ptr:
		var valid bool
		d.decode(reflect.ValueOf(&valid).Elem())
		// nil pointer, nothing to decode
		if !valid {
			return
		}
		// make sure we aren't decoding into nil
		if val.IsNil() {
			val.Set(reflect.New(val.Type().Elem()))
		}
		d.decode(val.Elem())
	case reflect.Bool:
		b := d.readN(1)
		if b[0] > 1 {
			panic("boolean value was not 0 or 1")
		}
		val.SetBool(b[0] == 1)
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		val.SetInt(DecInt64(d.readN(8)))
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		val.SetUint(DecUint64(d.readN(8)))
	case reflect.String:
		val.SetString(string(d.readPrefix()))
	case reflect.Slice:
		// slices are variable length, but otherwise the same as arrays.
		// just have to allocate them first, then we can fallthrough to the array logic.
		sliceLen := DecUint64(d.readN(8))
		// sanity-check the sliceLen, otherwise you can crash a peer by making
		// them allocate a massive slice
		if sliceLen > 1<<31-1 || sliceLen*uint64(val.Type().Elem().Size()) > maxSliceLen {
			panic("slice is too large")
		}
		val.Set(reflect.MakeSlice(val.Type(), int(sliceLen), int(sliceLen)))
		fallthrough
	case reflect.Array:
		// special case for byte arrays (e.g. hashes)
		if val.Type().Elem().Kind() == reflect.Uint8 {
			// convert val to a slice and read into it directly
			b := val.Slice(0, val.Len())
			_, err := io.ReadFull(d.r, b.Bytes())
			if err != nil {
				panic(err)
			}
			return
		}
		// arrays are unmarshalled by sequentially unmarshalling their elements
		for i := 0; i < val.Len(); i++ {
			d.decode(val.Index(i))
		}
		return
	case reflect.Struct:
		for i := 0; i < val.NumField(); i++ {
			d.decode(val.Field(i))
		}
		return
	default:
		panic("unknown type")
	}
}

// NewDecoder returns a new decoder that reads from r.
func NewDecoder(r io.Reader) *Decoder {
	return &Decoder{r}
}

// Unmarshal decodes the encoded value b and stores it in v, which must be a
// pointer. The decoding rules are the inverse of those specified in the
// package docstring.
func Unmarshal(b []byte, v interface{}) error {
	r := bytes.NewReader(b)
	return NewDecoder(r).Decode(v)
}

// ReadFile reads the contents of a file and decodes them into v.
func ReadFile(filename string, v interface{}) error {
	file, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer file.Close()
	return NewDecoder(file).Decode(v)
}
<commit_msg>extend length protection to strings<commit_after>// Package encoding converts arbitrary objects into byte slices, and vis
// versa. It also contains helper functions for reading and writing length-
// prefixed data. The encoding rules are as follows:
//
// Objects are encoded as binary data, without type information. The receiver
// uses context to determine the type to decode into.
//
// Integers are little-endian, and are always encoded as 8 bytes, i.e. their
// int64 or uint64 equivalent.
//
// Booleans are encoded as one byte, either zero (false) or one (true). No
// other values may be used.
//
// Nil pointers are equivalent to "false," i.e. a single zero byte. Valid
// pointers are represented by a "true" byte (0x01) followed by the encoding
// of the dereferenced value.
//
// Variable-length types, such as strings and slices, are represented by an 8-byte
// length-prefix followed by the encoded value.
//
// Slices and structs are simply the concatenation of their encoded elements.
// Byte slices are not subject to the 8-byte integer rule; they are encoded as
// their literal representation, one byte per byte.
//
// The ordering of struct fields is determined by their type definition. For
// example:
//
//   type foo struct { S string; I int }
//
//   Marshal(foo{"bar", 3}) == append(Marshal("bar"), Marshal(3)...)
//
// Finally, if a type implements the SiaMarshaler interface, its MarshalSia
// method will be used to encode the type. The resulting byte slice will be
// length-prefixed like any other variable-length type. During decoding, the
// type is decoded as a byte slice, and then passed to UnmarshalSia.
package encoding

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"os"
	"reflect"
)

// A SiaMarshaler can encode itself as a byte slice.
type SiaMarshaler interface {
	MarshalSia() []byte
}

// A SiaUnmarshaler can decode itself from a byte slice. If a decoding error
// occurs, UnmarshalSia should panic.
type SiaUnmarshaler interface {
	UnmarshalSia([]byte)
}

// An Encoder writes objects to an output stream.
type Encoder struct {
	w io.Writer
}

const (
	maxSliceLen = 4 * 1024 * 1024 // 4 MB
)

var (
	errBadPointer = errors.New("cannot decode into invalid pointer")
)

// Encode writes the encoding of v to the stream. For encoding details, see
// the package docstring.
func (e *Encoder) Encode(v interface{}) error {
	return e.encode(reflect.ValueOf(v))
}

// write catches instances where short writes do not return an error.
func (e *Encoder) write(p []byte) error {
	n, err := e.w.Write(p)
	if n != len(p) {
		return io.ErrShortWrite
	}
	return err
}

// Encode writes the encoding of val to the stream. For encoding details, see
// the package docstring.
func (e *Encoder) encode(val reflect.Value) error {
	// check for MarshalSia interface first
	if m, ok := val.Interface().(SiaMarshaler); ok {
		return WritePrefix(e.w, m.MarshalSia())
	}

	switch val.Kind() {
	case reflect.Ptr:
		// write either a 1 or 0
		if err := e.Encode(!val.IsNil()); err != nil {
			return err
		}
		if !val.IsNil() {
			return e.encode(val.Elem())
		}
	case reflect.Bool:
		if val.Bool() {
			return e.write([]byte{1})
		} else {
			return e.write([]byte{0})
		}
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		return e.write(EncInt64(val.Int()))
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		return e.write(EncUint64(val.Uint()))
	case reflect.String:
		return WritePrefix(e.w, []byte(val.String()))
	case reflect.Slice:
		// slices are variable length, so prepend the length and then fallthrough to array logic
		if err := e.write(EncUint64(uint64(val.Len()))); err != nil {
			return err
		}
		fallthrough
	case reflect.Array:
		// special case for byte arrays
		if val.Type().Elem().Kind() == reflect.Uint8 {
			// convert array to slice so we can use Bytes()
			// can't just use Slice() because array may be unaddressable
			slice := reflect.MakeSlice(reflect.SliceOf(val.Type().Elem()), val.Len(), val.Len())
			reflect.Copy(slice, val)
			return e.write(slice.Bytes())
		}
		// normal slices/arrays are encoded by sequentially encoding their elements
		for i := 0; i < val.Len(); i++ {
			if err := e.encode(val.Index(i)); err != nil {
				return err
			}
		}
	case reflect.Struct:
		for i := 0; i < val.NumField(); i++ {
			if err := e.encode(val.Field(i)); err != nil {
				return err
			}
		}
	default:
		// Marshalling should never fail. If it panics, you're doing something wrong,
		// like trying to encode a map or an unexported struct field.
		panic("could not marshal type " + val.Type().String())
	}
	return nil
}

// NewEncoder returns a new encoder that writes to w.
func NewEncoder(w io.Writer) *Encoder {
	return &Encoder{w}
}

// Marshal returns the encoding of v. For encoding details, see the package
// docstring.
func Marshal(v interface{}) []byte {
	b := new(bytes.Buffer)
	NewEncoder(b).Encode(v) // no error possible when using a bytes.Buffer
	return b.Bytes()
}

// MarshalAll encodes all of its inputs and returns their concatenation.
func MarshalAll(v ...interface{}) []byte {
	b := new(bytes.Buffer)
	enc := NewEncoder(b)
	for i := range v {
		enc.Encode(v[i]) // no error possible when using a bytes.Buffer
	}
	return b.Bytes()
}

// WriteFile writes v to a file. The file will be created if it does not exist.
func WriteFile(filename string, v interface{}) error {
	file, err := os.OpenFile(filename, os.O_WRONLY|os.O_CREATE, 0660)
	if err != nil {
		return err
	}
	defer file.Close()
	return NewEncoder(file).Encode(v)
}

// A Decoder reads and decodes values from an input stream.
type Decoder struct {
	r io.Reader
}

// Decode reads the next encoded value from its input stream and stores it in
// v, which must be a pointer. The decoding rules are the inverse of those
// specified in the package docstring.
func (d *Decoder) Decode(v interface{}) (err error) {
	// v must be a pointer
	pval := reflect.ValueOf(v)
	if pval.Kind() != reflect.Ptr || pval.IsNil() {
		return errBadPointer
	}

	// catch decoding panics and convert them to errors
	// note that this allows us to skip boundary checks during decoding
	defer func() {
		if r := recover(); r != nil {
			err = fmt.Errorf("could not decode type %s: %v", pval.Elem().Type().String(), r)
		}
	}()

	d.decode(pval.Elem())
	return
}

// readN reads n bytes and panics if the read fails.
func (d *Decoder) readN(n int) []byte {
	b := make([]byte, n)
	_, err := io.ReadFull(d.r, b)
	if err != nil {
		panic(err)
	}
	return b
}

// readPrefix reads a length-prefixed byte slice and panics if the read fails.
func (d *Decoder) readPrefix() []byte {
	// TODO: what should maxlen be?
	b, err := ReadPrefix(d.r, maxSliceLen)
	if err != nil {
		panic(err)
	}
	return b
}

// decode reads the next encoded value from its input stream and stores it in
// val. The decoding rules are the inverse of those specified in the package
// docstring.
func (d *Decoder) decode(val reflect.Value) {
	// check for UnmarshalSia interface first
	if val.CanAddr() {
		if u, ok := val.Addr().Interface().(SiaUnmarshaler); ok {
			u.UnmarshalSia(d.readPrefix())
			return
		}
	}

	switch val.Kind() {
	case reflect.Ptr:
		var valid bool
		d.decode(reflect.ValueOf(&valid).Elem())
		// nil pointer, nothing to decode
		if !valid {
			return
		}
		// make sure we aren't decoding into nil
		if val.IsNil() {
			val.Set(reflect.New(val.Type().Elem()))
		}
		d.decode(val.Elem())
	case reflect.Bool:
		b := d.readN(1)
		if b[0] > 1 {
			panic("boolean value was not 0 or 1")
		}
		val.SetBool(b[0] == 1)
	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
		val.SetInt(DecInt64(d.readN(8)))
	case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
		val.SetUint(DecUint64(d.readN(8)))
	case reflect.String:
		val.SetString(string(d.readPrefix()))
	case reflect.Slice:
		// slices are variable length, but otherwise the same as arrays.
		// just have to allocate them first, then we can fallthrough to the array logic.
		sliceLen := DecUint64(d.readN(8))
		// sanity-check the sliceLen, otherwise you can crash a peer by making
		// them allocate a massive slice
		if sliceLen > 1<<31-1 || sliceLen*uint64(val.Type().Elem().Size()) > maxSliceLen {
			panic("slice is too large")
		}
		val.Set(reflect.MakeSlice(val.Type(), int(sliceLen), int(sliceLen)))
		fallthrough
	case reflect.Array:
		// special case for byte arrays (e.g. hashes)
		if val.Type().Elem().Kind() == reflect.Uint8 {
			// convert val to a slice and read into it directly
			b := val.Slice(0, val.Len())
			_, err := io.ReadFull(d.r, b.Bytes())
			if err != nil {
				panic(err)
			}
			return
		}
		// arrays are unmarshalled by sequentially unmarshalling their elements
		for i := 0; i < val.Len(); i++ {
			d.decode(val.Index(i))
		}
		return
	case reflect.Struct:
		for i := 0; i < val.NumField(); i++ {
			d.decode(val.Field(i))
		}
		return
	default:
		panic("unknown type")
	}
}

// NewDecoder returns a new decoder that reads from r.
func NewDecoder(r io.Reader) *Decoder {
	return &Decoder{r}
}

// Unmarshal decodes the encoded value b and stores it in v, which must be a
// pointer. The decoding rules are the inverse of those specified in the
// package docstring.
func Unmarshal(b []byte, v interface{}) error {
	r := bytes.NewReader(b)
	return NewDecoder(r).Decode(v)
}

// ReadFile reads the contents of a file and decodes them into v.
func ReadFile(filename string, v interface{}) error {
	file, err := os.Open(filename)
	if err != nil {
		return err
	}
	defer file.Close()
	return NewDecoder(file).Decode(v)
}
<|endoftext|>


<commit_before>package core

import (
	"github.com/scrapli/scrapligo/driver/base"
	"github.com/scrapli/scrapligo/driver/network"
)

// NewJUNOSDriver return a driver setup for operation with Junos devices.
func NewJUNOSDriver(
	host string,
	options ...base.Option,
) (*network.Driver, error) {
	defaultPrivilegeLevels := map[string]*base.PrivilegeLevel{
		"exec": {
			Pattern:        `(?im)^({\w+:\d}\n){0,1}[\w\-@()/:]{1,63}>\s?$`,
			Name:           execPrivLevel,
			PreviousPriv:   "",
			Deescalate:     "",
			Escalate:       "",
			EscalateAuth:   false,
			EscalatePrompt: "",
		},
		"configuration": {
			Pattern:        `(?im)^({\w+:\d}\[edit\]\n){0,1}[\w\-@()/:]{1,63}#\s?$`,
			Name:           configPrivLevel,
			PreviousPriv:   execPrivLevel,
			Deescalate:     "exit configuration-mode",
			Escalate:       "configure",
			EscalateAuth:   false,
			EscalatePrompt: "",
		},
		"configuration_exclusive": {
			Pattern:        `(?im)^({\w+:\d}\[edit\]\n){0,1}[\w\-@()/:]{1,63}#\s?$`,
			Name:           "configuration_exclusive",
			PreviousPriv:   execPrivLevel,
			Deescalate:     "exit configuration-mode",
			Escalate:       "configure exclusive",
			EscalateAuth:   false,
			EscalatePrompt: "",
		},
		"configuration_private": {
			Pattern:        `(?im)^({\w+:\d}\[edit\]\n){0,1}[\w\-@()/:]{1,63}#\s?$`,
			Name:           "configuration_private",
			PreviousPriv:   execPrivLevel,
			Deescalate:     "exit configuration-mode",
			Escalate:       "configure exclusive",
			EscalateAuth:   false,
			EscalatePrompt: "",
		},
		"shell": {
			Pattern:            `(?im)^.*[%$]\s?$`,
			PatternNotContains: []string{"root"},
			Name:               "shell",
			PreviousPriv:       execPrivLevel,
			Deescalate:         "exit",
			Escalate:           "start shell",
			EscalateAuth:       false,
			EscalatePrompt:     "",
		},
		"root_shell": {
			Pattern:        `(?im)^.*root@[[:ascii:]]*:[[:ascii:]]*[%#]\s?$`,
			Name:           "root_shell",
			PreviousPriv:   execPrivLevel,
			Deescalate:     "exit",
			Escalate:       "start shell user root",
			EscalateAuth:   true,
			EscalatePrompt: `(?im)^[pP]assword:\s?$`,
		},
	}

	defaultFailedWhenContains := []string{
		"is ambiguous",
		"No valid completions",
		"unknown command",
		"syntax error",
	}

	const defaultDefaultDesiredPriv = execPrivLevel

	d, err := network.NewNetworkDriver(
		host,
		defaultPrivilegeLevels,
		defaultDefaultDesiredPriv,
		defaultFailedWhenContains,
		JUNOSOnOpen,
		JUNOSOnClose,
		options...)

	if err != nil {
		return nil, err
	}

	d.Augments["abortConfig"] = JUNOSAbortConfig

	return d, nil
}

// JUNOSOnOpen default on open callable for Junos.
func JUNOSOnOpen(d *network.Driver) error {
	err := d.AcquirePriv(d.DefaultDesiredPriv)
	if err != nil {
		return err
	}

	_, err = d.SendCommand("set cli screen-length 0", nil)
	if err != nil {
		return err
	}

	_, err = d.SendCommand("set cli screen-width 511", nil)
	if err != nil {
		return err
	}

	_, err = d.SendCommand("set cli complete-on-space off", nil)
	if err != nil {
		return err
	}

	return nil
}

// JUNOSOnClose default on close callable for Junos.
func JUNOSOnClose(d *network.Driver) error {
	err := d.AcquirePriv(d.DefaultDesiredPriv)
	if err != nil {
		return err
	}

	err = d.Channel.Write([]byte("exit"), false)
	if err != nil {
		return err
	}

	err = d.Channel.SendReturn()
	if err != nil {
		return err
	}

	return nil
}

// JUNOSAbortConfig abort Junos configuration session.
func JUNOSAbortConfig(d *network.Driver) (*base.Response, error) {
	_, _ = d.Channel.SendInput("rollback 0", false, false, -1)
	_, err := d.Channel.SendInput("exit", false, false, -1)
	d.CurrentPriv = privExecPrivLevel

	return nil, err
}
<commit_msg>expand junos root shell pattern<commit_after>package core

import (
	"github.com/scrapli/scrapligo/driver/base"
	"github.com/scrapli/scrapligo/driver/network"
)

// NewJUNOSDriver return a driver setup for operation with Junos devices.
func NewJUNOSDriver(
	host string,
	options ...base.Option,
) (*network.Driver, error) {
	defaultPrivilegeLevels := map[string]*base.PrivilegeLevel{
		"exec": {
			Pattern:        `(?im)^({\w+:\d}\n){0,1}[\w\-@()/:]{1,63}>\s?$`,
			Name:           execPrivLevel,
			PreviousPriv:   "",
			Deescalate:     "",
			Escalate:       "",
			EscalateAuth:   false,
			EscalatePrompt: "",
		},
		"configuration": {
			Pattern:        `(?im)^({\w+:\d}\[edit\]\n){0,1}[\w\-@()/:]{1,63}#\s?$`,
			Name:           configPrivLevel,
			PreviousPriv:   execPrivLevel,
			Deescalate:     "exit configuration-mode",
			Escalate:       "configure",
			EscalateAuth:   false,
			EscalatePrompt: "",
		},
		"configuration_exclusive": {
			Pattern:        `(?im)^({\w+:\d}\[edit\]\n){0,1}[\w\-@()/:]{1,63}#\s?$`,
			Name:           "configuration_exclusive",
			PreviousPriv:   execPrivLevel,
			Deescalate:     "exit configuration-mode",
			Escalate:       "configure exclusive",
			EscalateAuth:   false,
			EscalatePrompt: "",
		},
		"configuration_private": {
			Pattern:        `(?im)^({\w+:\d}\[edit\]\n){0,1}[\w\-@()/:]{1,63}#\s?$`,
			Name:           "configuration_private",
			PreviousPriv:   execPrivLevel,
			Deescalate:     "exit configuration-mode",
			Escalate:       "configure exclusive",
			EscalateAuth:   false,
			EscalatePrompt: "",
		},
		"shell": {
			Pattern:            `(?im)^.*[%$]\s?$`,
			PatternNotContains: []string{"root"},
			Name:               "shell",
			PreviousPriv:       execPrivLevel,
			Deescalate:         "exit",
			Escalate:           "start shell",
			EscalateAuth:       false,
			EscalatePrompt:     "",
		},
		"root_shell": {
			Pattern:        `(?im)^.*root@[[:ascii:]]*?:?[[:ascii:]]*?[%#]\s?$`,
			Name:           "root_shell",
			PreviousPriv:   execPrivLevel,
			Deescalate:     "exit",
			Escalate:       "start shell user root",
			EscalateAuth:   true,
			EscalatePrompt: `(?im)^[pP]assword:\s?$`,
		},
	}

	defaultFailedWhenContains := []string{
		"is ambiguous",
		"No valid completions",
		"unknown command",
		"syntax error",
	}

	const defaultDefaultDesiredPriv = execPrivLevel

	d, err := network.NewNetworkDriver(
		host,
		defaultPrivilegeLevels,
		defaultDefaultDesiredPriv,
		defaultFailedWhenContains,
		JUNOSOnOpen,
		JUNOSOnClose,
		options...)

	if err != nil {
		return nil, err
	}

	d.Augments["abortConfig"] = JUNOSAbortConfig

	return d, nil
}

// JUNOSOnOpen default on open callable for Junos.
func JUNOSOnOpen(d *network.Driver) error {
	err := d.AcquirePriv(d.DefaultDesiredPriv)
	if err != nil {
		return err
	}

	_, err = d.SendCommand("set cli screen-length 0", nil)
	if err != nil {
		return err
	}

	_, err = d.SendCommand("set cli screen-width 511", nil)
	if err != nil {
		return err
	}

	_, err = d.SendCommand("set cli complete-on-space off", nil)
	if err != nil {
		return err
	}

	return nil
}

// JUNOSOnClose default on close callable for Junos.
func JUNOSOnClose(d *network.Driver) error {
	err := d.AcquirePriv(d.DefaultDesiredPriv)
	if err != nil {
		return err
	}

	err = d.Channel.Write([]byte("exit"), false)
	if err != nil {
		return err
	}

	err = d.Channel.SendReturn()
	if err != nil {
		return err
	}

	return nil
}

// JUNOSAbortConfig abort Junos configuration session.
func JUNOSAbortConfig(d *network.Driver) (*base.Response, error) {
	_, _ = d.Channel.SendInput("rollback 0", false, false, -1)
	_, err := d.Channel.SendInput("exit", false, false, -1)
	d.CurrentPriv = privExecPrivLevel

	return nil, err
}
<|endoftext|>


<commit_before>package engine

import (
	"fmt"
	dockerclient "github.com/fsouza/go-dockerclient"
	log "github.com/jbdalido/logrus"
	"github.com/jbdalido/smg/utils"
	"io"
	"math/rand"
	"runtime"
	"strconv"
	"strings"
	"time"
)

type Container struct {
	ID               string
	Image            ImageName
	Name             string
	Tags             []string
	Hostname         string
	Code             int
	Privileged       bool
	Protection       bool
	Client           *dockerclient.Client
	HostConfig       *dockerclient.HostConfig
	ContainerConfig  *dockerclient.Config
	Docker           *dockerclient.Container
	Ports            []dockerclient.Port
	WorkingDirectory string
	Links            []*Container
}

// Inspect get the container definition
// and auto-protect the container against
// deletion, to avoid killing or removing
// already existing containers
func (c *Container) Inspect(id string) error {
	cont, err := c.Client.InspectContainer(id)
	if err != nil {
		return err
	}

	// Get intels
	c.Docker = cont
	c.Name = c.Docker.Name

	return nil
}

// Exists and IsRunning are wrapper of inspect
// to get a precise state of the container
// Exists is used to kill an existing container
func (c *Container) Exists(id string) bool {
	err := c.Inspect(id)
	if err != nil {
		return false
	}
	return true
}

// Exists and IsRunning are wrapper of inspect
// to get a precise state of the container
// IsRunning is used to attach to it
func (c *Container) IsRunning(id string) bool {

	err := c.Inspect(id)
	if err != nil {
		return false
	}

	if !c.Docker.State.Running {
		return false
	}

	return true
}

// Start is starting the container and
// redirection its log to stdout
func (c *Container) Start(out bool) (int, error) {

	if !c.IsRunning(c.Docker.ID) {

		err := c.Client.StartContainer(c.Docker.ID, c.HostConfig)
		if err != nil {
			return -1, err
		}

		if out {

			c.Logs(c.Docker.ID, utils.StdPre, utils.StdPre)
			c.Code, err = c.Client.WaitContainer(c.Docker.ID)
			if err != nil {
				return -1, nil
			}

		}
	}

	return c.Code, nil
}

// Logs is calling dockerclient logs function
// into a goroutine so we're not blocker by it
func (c *Container) Logs(id string, out, err io.Writer) {
	go func() {

		c.Client.Logs(dockerclient.LogsOptions{
			Container:    id,
			OutputStream: out,
			ErrorStream:  err,
			Follow:       true,
			Stdout:       true,
			Stderr:       true,
		})

	}()
}

// SetProtection is taking protective mesure to
// disable smuggler deletion and stopping functions
func (c *Container) SetProtection(p bool) {
	c.Protection = p
}

// Stop is stopping the container
func (c *Container) Stop() error {
	if c.Protection {
		return nil
	}
	if c.IsRunning(c.Docker.ID) {
		opts := dockerclient.KillContainerOptions{ID: c.Docker.ID}
		err := c.Client.KillContainer(opts)
		log.Debugf("Killing %s", c.Image.Name)
		if err != nil {
			log.Debugf("ERROR killing %s, %s", c.Image, err)
		}
	}
	return nil
}

// Delete is deleting the container
func (c *Container) Delete(force bool) error {
	// Keepalive means protection against suppression
	if c.Protection && !force {
		return nil
	}

	if c.Exists(c.Docker.ID) {
		// short the id for prettier logs
		id := c.Docker.ID[:8]
		opts := dockerclient.RemoveContainerOptions{ID: id}

		// Remove the container with the id
		err := c.Client.RemoveContainer(opts)
		if err != nil {
			log.Debugf("ERROR removing %s, %s", c.Image, err)
		}
	}
	return nil
}

// If system limits are defined in the smuggler conf file
// we're setting them to the container
func (c *Container) SetSystemLimits(cpu int64, ram int64) error {
	c.ContainerConfig.CpuShares = cpu
	c.ContainerConfig.Memory = ram
	return nil
}

func (c *Container) CreateDockerContainer() (err error) {

	opts := dockerclient.CreateContainerOptions{
		Name:   c.Name,
		Config: c.ContainerConfig,
	}

	c.Docker, err = c.Client.CreateContainer(opts)
	if err != nil {
		return fmt.Errorf("%s %s", c.Image, err)
	}
	return nil
}

// Set
func (c *Container) SetContainerConfig(env []string, cmd []string) error {

	c.ContainerConfig = &dockerclient.Config{
		Hostname:     c.Hostname,
		Image:        c.Image.ToString(),
		AttachStdin:  true,
		AttachStdout: true,
		AttachStderr: true,
	}

	if len(env) > 0 {
		c.ContainerConfig.Env = env
	}

	if len(cmd) > 0 {
		c.ContainerConfig.WorkingDir = "/data"
		c.ContainerConfig.Cmd = []string{"/data/run.sh"}
	}

	return nil
}

func (c *Container) SetBinds(binds []string) error {

	if c.HostConfig == nil {
		c.HostConfig = &dockerclient.HostConfig{}
	}
	// Volumes are set for both smugglers and users
	for _, bind := range binds {
		b := utils.EnvResolver(bind)
		n := strings.Split(b, ":")
		if len(n) != 2 {
			log.Warningf("Malformed Volume %s", b)
			continue
		}
		// TODO - Use another hidden directory
		if strings.HasPrefix(n[1], "/data/") {
			log.Fatalf("/data is currently user by smuggler. This directory cannot be mounted")
		}

		if runtime.GOOS == "darwin" {
			if !strings.HasPrefix(n[0], "/Users") {
				log.Warnf("Volumes will not work if outside /Users with boot2docker (%s)", n[0])
			}
		}

		log.Debugf("Binding volumes %s", b)
		c.HostConfig.Binds = append(c.HostConfig.Binds, b)
	}

	return nil

}

func (c *Container) SetPorts(ports []string) error {

	if c.HostConfig == nil {
		c.HostConfig = &dockerclient.HostConfig{}
	}

	rand.Seed(time.Now().UTC().UnixNano())

	// Set Host Port Binding
	bindPorts := map[dockerclient.Port][]dockerclient.PortBinding{}
	if len(ports) > 0 {
		for _, port := range ports {
			ps := strings.Split(string(port), ":")

			var (
				hostPort string
				p        string
			)
			// TODO:
			// Set the test against the random port, set a generate
			// function to obtain a free port
			if len(ps) == 2 {
				if utils.PortUsable(c.ParsePort(ps[0])) {
					hostPort = ps[0]
				} else {
					hostPort = strconv.Itoa(rand.Intn(10000) + 30000)
				}
				p = ps[1]
			} else {
				hostPort = strconv.Itoa(rand.Intn(10000) + 30000)
				p = port
			}

			bindPorts[c.ContainerPort(p)] = []dockerclient.PortBinding{{
				HostIp:   "0.0.0.0",
				HostPort: c.ParsePort(hostPort),
			}}
			log.Infof("Opening ports %s:%s", hostPort, p)
		}
	}

	c.HostConfig.PortBindings = bindPorts

	if c.Privileged {
		c.HostConfig.Privileged = true
	}

	return nil
}

func (c *Container) ParsePort(p string) string {
	// Let's check if a protocol is asked
	lenp := strings.Split(p, "/")
	if len(lenp) > 1 {
		return lenp[0]
	}
	return p

}

func (c *Container) ContainerPort(p string) dockerclient.Port {

	// Let's validate the protocol used
	lenp := strings.Split(p, "/")

	// A protocol is set, validation
	if len(lenp) == 2 {
		if lenp[1] == "tcp" || lenp[1] == "udp" {
			return dockerclient.Port(lenp[0] + "/" + lenp[1])
		}
		return dockerclient.Port(lenp[0] + "/tcp")
	}

	// No protocol, asign tcp
	return dockerclient.Port(p + "/tcp")
}

func (c *Container) addLinks(links []*Container) {

	var linked []string

	if len(links) > 0 {
		for _, l := range links {
			linked = append(linked, l.Name+":"+l.Hostname)
			log.Infof("Linking %s", l.Hostname)
		}
		c.HostConfig.Links = linked
	}
}
<commit_msg>Use the correct names<commit_after>package engine

import (
	"fmt"
	dockerclient "github.com/fsouza/go-dockerclient"
	log "github.com/jbdalido/logrus"
	"github.com/jbdalido/smg/utils"
	"io"
	"math/rand"
	"runtime"
	"strconv"
	"strings"
	"time"
)

type Container struct {
	ID               string
	Image            ImageName
	Name             string
	Tags             []string
	Hostname         string
	Code             int
	Privileged       bool
	Protection       bool
	Client           *dockerclient.Client
	HostConfig       *dockerclient.HostConfig
	ContainerConfig  *dockerclient.Config
	Docker           *dockerclient.Container
	Ports            []dockerclient.Port
	WorkingDirectory string
	Links            []*Container
}

// Inspect get the container definition
// and auto-protect the container against
// deletion, to avoid killing or removing
// already existing containers
func (c *Container) Inspect(id string) error {
	cont, err := c.Client.InspectContainer(id)
	if err != nil {
		return err
	}

	// Get intels
	c.Docker = cont
	c.Name = c.Docker.Name

	return nil
}

// Exists and IsRunning are wrapper of inspect
// to get a precise state of the container
// Exists is used to kill an existing container
func (c *Container) Exists(id string) bool {
	err := c.Inspect(id)
	if err != nil {
		return false
	}
	return true
}

// Exists and IsRunning are wrapper of inspect
// to get a precise state of the container
// IsRunning is used to attach to it
func (c *Container) IsRunning(id string) bool {

	err := c.Inspect(id)
	if err != nil {
		return false
	}

	if !c.Docker.State.Running {
		return false
	}

	return true
}

// Start is starting the container and
// redirection its log to stdout
func (c *Container) Start(out bool) (int, error) {

	if !c.IsRunning(c.Docker.ID) {

		err := c.Client.StartContainer(c.Docker.ID, c.HostConfig)
		if err != nil {
			return -1, err
		}

		if out {

			c.Logs(c.Docker.ID, utils.StdPre, utils.StdPre)
			c.Code, err = c.Client.WaitContainer(c.Docker.ID)
			if err != nil {
				return -1, nil
			}

		}
	}

	return c.Code, nil
}

// Logs is calling dockerclient logs function
// into a goroutine so we're not blocker by it
func (c *Container) Logs(id string, out, err io.Writer) {
	go func() {

		c.Client.Logs(dockerclient.LogsOptions{
			Container:    id,
			OutputStream: out,
			ErrorStream:  err,
			Follow:       true,
			Stdout:       true,
			Stderr:       true,
		})

	}()
}

// SetProtection is taking protective mesure to
// disable smuggler deletion and stopping functions
func (c *Container) SetProtection(p bool) {
	c.Protection = p
}

// Stop is stopping the container
func (c *Container) Stop() error {
	if c.Protection {
		return nil
	}
	if c.IsRunning(c.Docker.ID) {
		opts := dockerclient.KillContainerOptions{ID: c.Docker.ID}
		err := c.Client.KillContainer(opts)
		log.Debugf("Killing %s", c.Image.Name)
		if err != nil {
			log.Debugf("ERROR killing %s, %s", c.Image, err)
		}
	}
	return nil
}

// Delete is deleting the container
func (c *Container) Delete(force bool) error {
	// Keepalive means protection against suppression
	if c.Protection && !force {
		return nil
	}

	if c.Exists(c.Docker.ID) {
		// short the id for prettier logs
		id := c.Docker.ID[:8]
		opts := dockerclient.RemoveContainerOptions{ID: id}

		// Remove the container with the id
		err := c.Client.RemoveContainer(opts)
		if err != nil {
			log.Debugf("ERROR removing %s, %s", c.Image, err)
		}
	}
	return nil
}

// If system limits are defined in the smuggler conf file
// we're setting them to the container
func (c *Container) SetSystemLimits(cpu int64, ram int64) error {
	c.ContainerConfig.CPUShares = cpu
	c.ContainerConfig.Memory = ram
	return nil
}

func (c *Container) CreateDockerContainer() (err error) {

	opts := dockerclient.CreateContainerOptions{
		Name:   c.Name,
		Config: c.ContainerConfig,
	}

	c.Docker, err = c.Client.CreateContainer(opts)
	if err != nil {
		return fmt.Errorf("%s %s", c.Image, err)
	}
	return nil
}

// Set
func (c *Container) SetContainerConfig(env []string, cmd []string) error {

	c.ContainerConfig = &dockerclient.Config{
		Hostname:     c.Hostname,
		Image:        c.Image.ToString(),
		AttachStdin:  true,
		AttachStdout: true,
		AttachStderr: true,
	}

	if len(env) > 0 {
		c.ContainerConfig.Env = env
	}

	if len(cmd) > 0 {
		c.ContainerConfig.WorkingDir = "/data"
		c.ContainerConfig.Cmd = []string{"/data/run.sh"}
	}

	return nil
}

func (c *Container) SetBinds(binds []string) error {

	if c.HostConfig == nil {
		c.HostConfig = &dockerclient.HostConfig{}
	}
	// Volumes are set for both smugglers and users
	for _, bind := range binds {
		b := utils.EnvResolver(bind)
		n := strings.Split(b, ":")
		if len(n) != 2 {
			log.Warningf("Malformed Volume %s", b)
			continue
		}
		// TODO - Use another hidden directory
		if strings.HasPrefix(n[1], "/data/") {
			log.Fatalf("/data is currently user by smuggler. This directory cannot be mounted")
		}

		if runtime.GOOS == "darwin" {
			if !strings.HasPrefix(n[0], "/Users") {
				log.Warnf("Volumes will not work if outside /Users with boot2docker (%s)", n[0])
			}
		}

		log.Debugf("Binding volumes %s", b)
		c.HostConfig.Binds = append(c.HostConfig.Binds, b)
	}

	return nil

}

func (c *Container) SetPorts(ports []string) error {

	if c.HostConfig == nil {
		c.HostConfig = &dockerclient.HostConfig{}
	}

	rand.Seed(time.Now().UTC().UnixNano())

	// Set Host Port Binding
	bindPorts := map[dockerclient.Port][]dockerclient.PortBinding{}
	if len(ports) > 0 {
		for _, port := range ports {
			ps := strings.Split(string(port), ":")

			var (
				hostPort string
				p        string
			)
			// TODO:
			// Set the test against the random port, set a generate
			// function to obtain a free port
			if len(ps) == 2 {
				if utils.PortUsable(c.ParsePort(ps[0])) {
					hostPort = ps[0]
				} else {
					hostPort = strconv.Itoa(rand.Intn(10000) + 30000)
				}
				p = ps[1]
			} else {
				hostPort = strconv.Itoa(rand.Intn(10000) + 30000)
				p = port
			}

			bindPorts[c.ContainerPort(p)] = []dockerclient.PortBinding{{
				HostIP:   "0.0.0.0",
				HostPort: c.ParsePort(hostPort),
			}}
			log.Infof("Opening ports %s:%s", hostPort, p)
		}
	}

	c.HostConfig.PortBindings = bindPorts

	if c.Privileged {
		c.HostConfig.Privileged = true
	}

	return nil
}

func (c *Container) ParsePort(p string) string {
	// Let's check if a protocol is asked
	lenp := strings.Split(p, "/")
	if len(lenp) > 1 {
		return lenp[0]
	}
	return p

}

func (c *Container) ContainerPort(p string) dockerclient.Port {

	// Let's validate the protocol used
	lenp := strings.Split(p, "/")

	// A protocol is set, validation
	if len(lenp) == 2 {
		if lenp[1] == "tcp" || lenp[1] == "udp" {
			return dockerclient.Port(lenp[0] + "/" + lenp[1])
		}
		return dockerclient.Port(lenp[0] + "/tcp")
	}

	// No protocol, asign tcp
	return dockerclient.Port(p + "/tcp")
}

func (c *Container) addLinks(links []*Container) {

	var linked []string

	if len(links) > 0 {
		for _, l := range links {
			linked = append(linked, l.Name+":"+l.Hostname)
			log.Infof("Linking %s", l.Hostname)
		}
		c.HostConfig.Links = linked
	}
}
<|endoftext|>


<commit_before>package clients

import (
	"container/list"
	"sync"
	"time"

	"github.com/VolantMQ/persistence"
	"github.com/VolantMQ/volantmq/connection"
	"github.com/VolantMQ/volantmq/packet"
	"github.com/VolantMQ/volantmq/subscriber"
	"github.com/VolantMQ/volantmq/topics/types"
	"github.com/VolantMQ/volantmq/types"
	"go.uber.org/zap"
)

type sessionEvents interface {
	sessionOffline(string, bool, *expiry)
	connectionClosed(string, packet.ReasonCode)
	subscriberShutdown(string, subscriber.SessionProvider)
}

type sessionPreConfig struct {
	id          string
	createdAt   time.Time
	messenger   types.TopicMessenger
	conn        connection.Session
	persistence persistence.Packets
}

type sessionConfig struct {
	sessionEvents
	subscriber subscriber.SessionProvider
	will       *packet.Publish
	expireIn   *uint32
	willDelay  uint32
	durable    bool
	version    packet.ProtocolVersion
}

type session struct {
	sessionPreConfig
	log     *zap.Logger
	idLock  *sync.Mutex
	lock    sync.Mutex
	stopReq types.Once
	sessionConfig
}

type temporaryPublish struct {
	gList *list.List
	qList *list.List
}

func newTmpPublish() *temporaryPublish {
	return &temporaryPublish{
		gList: list.New(),
		qList: list.New(),
	}
}

func (t *temporaryPublish) Publish(id string, p *packet.Publish) {
	if p.QoS() == packet.QoS0 {
		t.gList.PushBack(p)
	} else {
		t.qList.PushBack(p)
	}
}

func newSession(c sessionPreConfig) *session {
	s := &session{
		sessionPreConfig: c,
	}

	return s
}

func (s *session) configure(c sessionConfig, clean bool) {
	s.sessionConfig = c

	s.conn.SetOptions(connection.AttachSession(s))

	if !clean {
		tmp := newTmpPublish()
		s.subscriber.Online(tmp)
		s.persistence.PacketsForEach([]byte(s.id), s.conn)
		s.subscriber.Online(s.conn)
		s.conn.LoadRemaining(tmp.gList, tmp.qList)
	} else {
		s.subscriber.Online(s.conn)
	}
}

func (s *session) start() {
	s.idLock.Unlock()
}

func (s *session) stop(reason packet.ReasonCode) {
	s.stopReq.Do(func() {
		s.conn.Stop(reason)
	})

	//s.wgDisconnected.Wait()

	//state := &persistence.SessionState{
	//	Timestamp: s.createdAt.Format(time.RFC3339),
	//}

	//if s.expireIn != nil || (s.willDelay > 0 && s.will != nil) {
	//	state.Expire = &persistence.SessionDelays{
	//		Since: s.expiringSince.Format(time.RFC3339),
	//	}
	//
	//	elapsed := uint32(time.Since(s.expiringSince) / time.Second)
	//
	//	if (s.willDelay > 0 && s.will != nil) && (s.willDelay-elapsed) > 0 {
	//		s.willDelay = s.willDelay - elapsed
	//		s.will.SetPacketID(0)
	//		if buf, err := packet.Encode(s.will); err != nil {
	//
	//		} else {
	//			state.Expire.WillIn = strconv.Itoa(int(s.willDelay))
	//			state.Expire.WillData = buf
	//		}
	//	}
	//
	//	if s.expireIn != nil && *s.expireIn > 0 && (*s.expireIn-elapsed) > 0 {
	//		*s.expireIn = *s.expireIn - elapsed
	//	}
	//}

	//return state
}

// SignalPublish process PUBLISH packet from client
func (s *session) SignalPublish(pkt *packet.Publish) error {
	pkt.SetPublishID(s.subscriber.Hash())

	// [MQTT-3.3.1.3]
	if pkt.Retain() {
		if err := s.messenger.Retain(pkt); err != nil {
			s.log.Error("Error retaining message", zap.String("ClientID", s.id), zap.Error(err))
		}

		// [MQTT-3.3.1-7]
		if pkt.QoS() == packet.QoS0 {
			retained := packet.NewPublish(s.version)
			retained.SetQoS(pkt.QoS())     // nolint: errcheck
			retained.SetTopic(pkt.Topic()) // nolint: errcheck
			//s.retained.list = append(s.retained.list, m)
		}
	}

	if err := s.messenger.Publish(pkt); err != nil {
		s.log.Error("Couldn't publish", zap.String("ClientID", s.id), zap.Error(err))
	}

	return nil
}

// SignalSubscribe process SUBSCRIBE packet from client
func (s *session) SignalSubscribe(pkt *packet.Subscribe) (packet.Provider, error) {
	m, _ := packet.New(s.version, packet.SUBACK)
	resp, _ := m.(*packet.SubAck)

	id, _ := pkt.ID()
	resp.SetPacketID(id)

	var retCodes []packet.ReasonCode
	var retainedPublishes []*packet.Publish

	pkt.RangeTopics(func(t string, ops packet.SubscriptionOptions) {
		reason := packet.CodeSuccess // nolint: ineffassign
		//authorized := true
		// TODO: check permissions here

		//if authorized {
		subsID := uint32(0)

		// V5.0 [MQTT-3.8.2.1.2]
		if prop := pkt.PropertyGet(packet.PropertySubscriptionIdentifier); prop != nil {
			if v, e := prop.AsInt(); e == nil {
				subsID = v
			}
		}

		subsParams := topicsTypes.SubscriptionParams{
			ID:  subsID,
			Ops: ops,
		}

		if grantedQoS, retained, err := s.subscriber.Subscribe(t, &subsParams); err != nil {
			// [MQTT-3.9.3]
			if s.version == packet.ProtocolV50 {
				reason = packet.CodeUnspecifiedError
			} else {
				reason = packet.QosFailure
			}
		} else {
			reason = packet.ReasonCode(grantedQoS)
			retainedPublishes = append(retainedPublishes, retained...)
		}

		retCodes = append(retCodes, reason)
	})

	if err := resp.AddReturnCodes(retCodes); err != nil {
		return nil, err
	}

	// Now put retained messages into publish queue
	for _, rp := range retainedPublishes {
		if pkt, err := rp.Clone(s.version); err == nil {
			pkt.SetRetain(true)
			s.conn.Publish(s.id, pkt)
		} else {
			s.log.Error("Couldn't clone PUBLISH message", zap.String("ClientID", s.id), zap.Error(err))
		}
	}

	return resp, nil
}

// SignalUnSubscribe process UNSUBSCRIBE packet from client
func (s *session) SignalUnSubscribe(pkt *packet.UnSubscribe) (packet.Provider, error) {
	var retCodes []packet.ReasonCode

	for _, t := range pkt.Topics() {
		// TODO: check permissions here
		authorized := true
		reason := packet.CodeSuccess

		if authorized {
			if err := s.subscriber.UnSubscribe(t); err != nil {
				s.log.Error("Couldn't unsubscribe from topic", zap.Error(err))
				reason = packet.CodeNoSubscriptionExisted
			}
		} else {
			reason = packet.CodeNotAuthorized
		}

		retCodes = append(retCodes, reason)
	}

	m, _ := packet.New(s.version, packet.UNSUBACK)
	resp, _ := m.(*packet.UnSubAck)

	id, _ := pkt.ID()
	resp.SetPacketID(id)
	resp.AddReturnCodes(retCodes) // nolint: errcheck

	return resp, nil
}

// SignalDisconnect process DISCONNECT packet from client
func (s *session) SignalDisconnect(pkt *packet.Disconnect) (packet.Provider, error) {
	var err error

	err = packet.CodeSuccess

	if s.version == packet.ProtocolV50 {
		// FIXME: CodeRefusedBadUsernameOrPassword has same id as CodeDisconnectWithWill
		if pkt.ReasonCode() != packet.CodeRefusedBadUsernameOrPassword {
			s.will = nil
		}

		if prop := pkt.PropertyGet(packet.PropertySessionExpiryInterval); prop != nil {
			if val, ok := prop.AsInt(); ok == nil {
				// If the Session Expiry Interval in the CONNECT packet was zero, then it is a Protocol Error to set a non-
				// zero Session Expiry Interval in the DISCONNECT packet sent by the Client. If such a non-zero Session
				// Expiry Interval is received by the Server, it does not treat it as a valid DISCONNECT packet. The Server
				// uses DISCONNECT with Reason Code 0x82 (Protocol Error) as described in section 4.13.
				if (s.expireIn != nil && *s.expireIn == 0) && val != 0 {
					err = packet.CodeProtocolError
				} else {
					s.expireIn = &val
				}
			}
		}
	} else {
		s.will = nil
	}

	return nil, err
}

// SignalOffline put subscriber in offline mode
func (s *session) SignalOffline() {
	s.subscriber.Offline(!s.durable)
}

// SignalConnectionClose net connection has been closed
func (s *session) SignalConnectionClose(params connection.DisconnectParams) {
	// If session expiry is set to 0, the Session ends when the Network Connection is closed
	if s.expireIn != nil && *s.expireIn == 0 {
		s.durable = true
	}

	// valid willMsg pointer tells we have will message
	// if session is clean send will regardless to will delay
	if s.will != nil && s.willDelay == 0 {
		if err := s.messenger.Publish(s.will); err != nil {
			s.log.Error("Publish will message", zap.String("ClientID", s.id), zap.Error(err))
		}
		s.will = nil
	}

	s.connectionClosed(s.id, params.Reason)

	if s.durable && len(params.Packets) > 0 {
		s.persistence.PacketsStore([]byte(s.id), params.Packets)
	}

	keepContainer := s.durable && s.subscriber.HasSubscriptions()

	if !s.durable || !s.subscriber.HasSubscriptions() {
		s.subscriberShutdown(s.id, s.subscriber)
		s.subscriber = nil
	}

	var exp *expiry

	if params.Reason != packet.CodeSessionTakenOver {
		if s.willDelay > 0 || (s.expireIn != nil && *s.expireIn > 0) {
			exp = newExpiry(
				expiryConfig{
					id:        s.id,
					createdAt: s.createdAt,
					messenger: s.messenger,
					will:      s.will,
					expireIn:  s.expireIn,
					willDelay: s.willDelay,
				})

			keepContainer = true
		}
	}

	s.sessionOffline(s.id, keepContainer, exp)

	s.stopReq.Do(func() {})
	s.conn = nil
}
<commit_msg>Init session log when object is created<commit_after>package clients

import (
	"container/list"
	"sync"
	"time"

	"github.com/VolantMQ/persistence"
	"github.com/VolantMQ/volantmq/configuration"
	"github.com/VolantMQ/volantmq/connection"
	"github.com/VolantMQ/volantmq/packet"
	"github.com/VolantMQ/volantmq/subscriber"
	"github.com/VolantMQ/volantmq/topics/types"
	"github.com/VolantMQ/volantmq/types"
	"go.uber.org/zap"
)

type sessionEvents interface {
	sessionOffline(string, bool, *expiry)
	connectionClosed(string, packet.ReasonCode)
	subscriberShutdown(string, subscriber.SessionProvider)
}

type sessionPreConfig struct {
	id          string
	createdAt   time.Time
	messenger   types.TopicMessenger
	conn        connection.Session
	persistence persistence.Packets
}

type sessionConfig struct {
	sessionEvents
	subscriber subscriber.SessionProvider
	will       *packet.Publish
	expireIn   *uint32
	willDelay  uint32
	durable    bool
	version    packet.ProtocolVersion
}

type session struct {
	sessionPreConfig
	log     *zap.Logger
	idLock  *sync.Mutex
	lock    sync.Mutex
	stopReq types.Once
	sessionConfig
}

type temporaryPublish struct {
	gList *list.List
	qList *list.List
}

func newTmpPublish() *temporaryPublish {
	return &temporaryPublish{
		gList: list.New(),
		qList: list.New(),
	}
}

func (t *temporaryPublish) Publish(id string, p *packet.Publish) {
	if p.QoS() == packet.QoS0 {
		t.gList.PushBack(p)
	} else {
		t.qList.PushBack(p)
	}
}

func newSession(c sessionPreConfig) *session {
	s := &session{
		sessionPreConfig: c,
		log:              configuration.GetLogger().Named("session"),
	}

	return s
}

func (s *session) configure(c sessionConfig, clean bool) {
	s.sessionConfig = c

	s.conn.SetOptions(connection.AttachSession(s))

	if !clean {
		tmp := newTmpPublish()
		s.subscriber.Online(tmp)
		s.persistence.PacketsForEach([]byte(s.id), s.conn)
		s.subscriber.Online(s.conn)
		s.conn.LoadRemaining(tmp.gList, tmp.qList)
	} else {
		s.subscriber.Online(s.conn)
	}
}

func (s *session) start() {
	s.idLock.Unlock()
}

func (s *session) stop(reason packet.ReasonCode) {
	s.stopReq.Do(func() {
		s.conn.Stop(reason)
	})

	//s.wgDisconnected.Wait()

	//state := &persistence.SessionState{
	//	Timestamp: s.createdAt.Format(time.RFC3339),
	//}

	//if s.expireIn != nil || (s.willDelay > 0 && s.will != nil) {
	//	state.Expire = &persistence.SessionDelays{
	//		Since: s.expiringSince.Format(time.RFC3339),
	//	}
	//
	//	elapsed := uint32(time.Since(s.expiringSince) / time.Second)
	//
	//	if (s.willDelay > 0 && s.will != nil) && (s.willDelay-elapsed) > 0 {
	//		s.willDelay = s.willDelay - elapsed
	//		s.will.SetPacketID(0)
	//		if buf, err := packet.Encode(s.will); err != nil {
	//
	//		} else {
	//			state.Expire.WillIn = strconv.Itoa(int(s.willDelay))
	//			state.Expire.WillData = buf
	//		}
	//	}
	//
	//	if s.expireIn != nil && *s.expireIn > 0 && (*s.expireIn-elapsed) > 0 {
	//		*s.expireIn = *s.expireIn - elapsed
	//	}
	//}

	//return state
}

// SignalPublish process PUBLISH packet from client
func (s *session) SignalPublish(pkt *packet.Publish) error {
	pkt.SetPublishID(s.subscriber.Hash())

	// [MQTT-3.3.1.3]
	if pkt.Retain() {
		if err := s.messenger.Retain(pkt); err != nil {
			s.log.Error("Error retaining message", zap.String("ClientID", s.id), zap.Error(err))
		}

		// [MQTT-3.3.1-7]
		if pkt.QoS() == packet.QoS0 {
			retained := packet.NewPublish(s.version)
			retained.SetQoS(pkt.QoS())     // nolint: errcheck
			retained.SetTopic(pkt.Topic()) // nolint: errcheck
			//s.retained.list = append(s.retained.list, m)
		}
	}

	if err := s.messenger.Publish(pkt); err != nil {
		s.log.Error("Couldn't publish", zap.String("ClientID", s.id), zap.Error(err))
	}

	return nil
}

// SignalSubscribe process SUBSCRIBE packet from client
func (s *session) SignalSubscribe(pkt *packet.Subscribe) (packet.Provider, error) {
	m, _ := packet.New(s.version, packet.SUBACK)
	resp, _ := m.(*packet.SubAck)

	id, _ := pkt.ID()
	resp.SetPacketID(id)

	var retCodes []packet.ReasonCode
	var retainedPublishes []*packet.Publish

	pkt.RangeTopics(func(t string, ops packet.SubscriptionOptions) {
		reason := packet.CodeSuccess // nolint: ineffassign
		//authorized := true
		// TODO: check permissions here

		//if authorized {
		subsID := uint32(0)

		// V5.0 [MQTT-3.8.2.1.2]
		if prop := pkt.PropertyGet(packet.PropertySubscriptionIdentifier); prop != nil {
			if v, e := prop.AsInt(); e == nil {
				subsID = v
			}
		}

		subsParams := topicsTypes.SubscriptionParams{
			ID:  subsID,
			Ops: ops,
		}

		if grantedQoS, retained, err := s.subscriber.Subscribe(t, &subsParams); err != nil {
			// [MQTT-3.9.3]
			if s.version == packet.ProtocolV50 {
				reason = packet.CodeUnspecifiedError
			} else {
				reason = packet.QosFailure
			}
		} else {
			reason = packet.ReasonCode(grantedQoS)
			retainedPublishes = append(retainedPublishes, retained...)
		}

		retCodes = append(retCodes, reason)
	})

	if err := resp.AddReturnCodes(retCodes); err != nil {
		return nil, err
	}

	// Now put retained messages into publish queue
	for _, rp := range retainedPublishes {
		if pkt, err := rp.Clone(s.version); err == nil {
			pkt.SetRetain(true)
			s.conn.Publish(s.id, pkt)
		} else {
			s.log.Error("Couldn't clone PUBLISH message", zap.String("ClientID", s.id), zap.Error(err))
		}
	}

	return resp, nil
}

// SignalUnSubscribe process UNSUBSCRIBE packet from client
func (s *session) SignalUnSubscribe(pkt *packet.UnSubscribe) (packet.Provider, error) {
	var retCodes []packet.ReasonCode

	for _, t := range pkt.Topics() {
		// TODO: check permissions here
		authorized := true
		reason := packet.CodeSuccess

		if authorized {
			if err := s.subscriber.UnSubscribe(t); err != nil {
				s.log.Error("Couldn't unsubscribe from topic", zap.Error(err))
				reason = packet.CodeNoSubscriptionExisted
			}
		} else {
			reason = packet.CodeNotAuthorized
		}

		retCodes = append(retCodes, reason)
	}

	m, _ := packet.New(s.version, packet.UNSUBACK)
	resp, _ := m.(*packet.UnSubAck)

	id, _ := pkt.ID()
	resp.SetPacketID(id)
	resp.AddReturnCodes(retCodes) // nolint: errcheck

	return resp, nil
}

// SignalDisconnect process DISCONNECT packet from client
func (s *session) SignalDisconnect(pkt *packet.Disconnect) (packet.Provider, error) {
	var err error

	err = packet.CodeSuccess

	if s.version == packet.ProtocolV50 {
		// FIXME: CodeRefusedBadUsernameOrPassword has same id as CodeDisconnectWithWill
		if pkt.ReasonCode() != packet.CodeRefusedBadUsernameOrPassword {
			s.will = nil
		}

		if prop := pkt.PropertyGet(packet.PropertySessionExpiryInterval); prop != nil {
			if val, ok := prop.AsInt(); ok == nil {
				// If the Session Expiry Interval in the CONNECT packet was zero, then it is a Protocol Error to set a non-
				// zero Session Expiry Interval in the DISCONNECT packet sent by the Client. If such a non-zero Session
				// Expiry Interval is received by the Server, it does not treat it as a valid DISCONNECT packet. The Server
				// uses DISCONNECT with Reason Code 0x82 (Protocol Error) as described in section 4.13.
				if (s.expireIn != nil && *s.expireIn == 0) && val != 0 {
					err = packet.CodeProtocolError
				} else {
					s.expireIn = &val
				}
			}
		}
	} else {
		s.will = nil
	}

	return nil, err
}

// SignalOffline put subscriber in offline mode
func (s *session) SignalOffline() {
	s.subscriber.Offline(!s.durable)
}

// SignalConnectionClose net connection has been closed
func (s *session) SignalConnectionClose(params connection.DisconnectParams) {
	// If session expiry is set to 0, the Session ends when the Network Connection is closed
	if s.expireIn != nil && *s.expireIn == 0 {
		s.durable = true
	}

	// valid willMsg pointer tells we have will message
	// if session is clean send will regardless to will delay
	if s.will != nil && s.willDelay == 0 {
		if err := s.messenger.Publish(s.will); err != nil {
			s.log.Error("Publish will message", zap.String("ClientID", s.id), zap.Error(err))
		}
		s.will = nil
	}

	s.connectionClosed(s.id, params.Reason)

	if s.durable && len(params.Packets) > 0 {
		s.persistence.PacketsStore([]byte(s.id), params.Packets)
	}

	keepContainer := s.durable && s.subscriber.HasSubscriptions()

	if !s.durable || !s.subscriber.HasSubscriptions() {
		s.subscriberShutdown(s.id, s.subscriber)
		s.subscriber = nil
	}

	var exp *expiry

	if params.Reason != packet.CodeSessionTakenOver {
		if s.willDelay > 0 || (s.expireIn != nil && *s.expireIn > 0) {
			exp = newExpiry(
				expiryConfig{
					id:        s.id,
					createdAt: s.createdAt,
					messenger: s.messenger,
					will:      s.will,
					expireIn:  s.expireIn,
					willDelay: s.willDelay,
				})

			keepContainer = true
		}
	}

	s.sessionOffline(s.id, keepContainer, exp)

	s.stopReq.Do(func() {})
	s.conn = nil
}
<|endoftext|>