<commit_before>package securepassctl

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
)

const (
	// DefaultRemote is the default Content-Type header used in HTTP requests
	DefaultRemote = "https://beta.secure-pass.net"
	// ContentType is the default Content-Type header used in HTTP requests
	ContentType = "application/json"
	// UserAgent contains the default User-Agent value used in HTTP requests
	UserAgent = "SecurePass CLI"
)

// DebugLogger collects all debug messages
var DebugLogger = log.New(ioutil.Discard, "", log.LstdFlags)

// SecurePass main object type
type SecurePass struct {
	AppID     string `ini:"APP_ID"`
	AppSecret string `ini:"APP_SECRET"`
	Endpoint  string
}

func (s *SecurePass) setupRequestFieds(req *http.Request) {
	req.Header.Set("Accept", ContentType)
	req.Header.Set("Content-Type", ContentType)
	req.Header.Set("User-Agent", UserAgent)
	req.Header.Set("X-SecurePass-App-ID", s.AppID)
	req.Header.Set("X-SecurePass-App-Secret", s.AppSecret)
}

func (s *SecurePass) makeRequestURL(path string) (string, error) {
	baseURL, _ := url.Parse(s.Endpoint)
	URL, err := url.Parse(path)
	if err != nil {
		return "", err
	}
	return baseURL.ResolveReference(URL).String(), nil
}

// NewRequest initializes and issues an HTTP request to the SecurePass endpoint
func (s *SecurePass) NewRequest(method, path string, data *url.Values) (*http.Request, error) {
	var err error
	var req *http.Request

	URL, err := s.makeRequestURL(path)
	if err != nil {
		return nil, err
	}
	if data != nil {
		req, err = http.NewRequest(method, URL, bytes.NewBufferString(data.Encode()))
	} else {
		req, err = http.NewRequest(method, URL, nil)
	}
	if err != nil {
		return nil, err
	}
	s.setupRequestFieds(req)
	return req, nil
}

// DoRequest issues an HTTP request
func (s *SecurePass) DoRequest(req *http.Request, obj APIResponse, expstatus int) error {
	client := NewClient(nil)
	DebugLogger.Printf("Sending %s request to %s", req.Method, req.URL)
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	if resp.StatusCode != expstatus {
		return fmt.Errorf("%s", resp.Status)
	}
	defer resp.Body.Close()

	err = json.NewDecoder(resp.Body).Decode(obj)
	if err != nil {
		return err
	}
	if obj.ErrorCode() != 0 {
		return fmt.Errorf("%d: %s", obj.ErrorCode(), obj.ErrorMessage())
	}
	return nil
}

// NewClient initialize http.Client with a certain http.Transport
func NewClient(tr *http.Transport) *http.Client {
	// Skip SSL certificate verification
	if tr == nil {
		tr = &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: true}}
	}
	return &http.Client{Transport: tr}
}

// AppInfo retrieves information on a SecurePass application
func (s *SecurePass) AppInfo(app string) (*AppInfoResponse, error) {
	var obj AppInfoResponse

	data := url.Values{}
	if app != "" {
		data.Set("APP_ID", app)
	}

	req, err := s.NewRequest("POST", "/api/v1/apps/info", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// AppAdd represents /api/v1/apps/add
func (s *SecurePass) AppAdd(app *ApplicationDescriptor) (*AppAddResponse, error) {
	var obj AppAddResponse

	data := url.Values{}
	data.Set("LABEL", app.Label)
	data.Set("WRITE", fmt.Sprintf("%v", app.Write))
	data.Set("PRIVACY", fmt.Sprintf("%v", app.Privacy))
	if app.AllowNetworkIPv4 != "" {
		data.Set("ALLOW_NETWORK_IPv4", app.AllowNetworkIPv4)
	}
	if app.AllowNetworkIPv6 != "" {
		data.Set("ALLOW_NETWORK_IPv6", app.AllowNetworkIPv6)
	}
	if app.Group != "" {
		data.Set("GROUP", app.Group)
	}
	if app.Realm != "" {
		data.Set("REALM", app.Realm)
	}

	req, err := s.NewRequest("POST", "/api/v1/apps/add", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// AppDel deletes an application from SecurePass
func (s *SecurePass) AppDel(app string) (*Response, error) {
	var obj Response

	data := url.Values{}
	if app != "" {
		data.Set("APP_ID", app)
	}

	req, err := s.NewRequest("POST", "/api/v1/apps/delete", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// AppMod represents /api/v1/apps/modify
func (s *SecurePass) AppMod(appID string, app *ApplicationDescriptor) (*AppInfoResponse, error) {
	var obj AppInfoResponse

	data := url.Values{}
	data.Set("APP_ID", appID)
	data.Set("WRITE", fmt.Sprintf("%v", app.Write))
	data.Set("PRIVACY", fmt.Sprintf("%v", app.Privacy))
	if app.Label != "" {
		data.Set("LABEL", app.Label)
	}
	if app.AllowNetworkIPv4 != "" {
		data.Set("ALLOW_NETWORK_IPv4", app.AllowNetworkIPv4)
	}
	if app.AllowNetworkIPv6 != "" {
		data.Set("ALLOW_NETWORK_IPv6", app.AllowNetworkIPv6)
	}
	if app.Group != "" {
		data.Set("GROUP", app.Group)
	}

	req, err := s.NewRequest("POST", "/api/v1/apps/modify", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// AppList retrieves the list of applications available in SecurePass
func (s *SecurePass) AppList(realm string) (*AppListResponse, error) {
	var obj AppListResponse

	data := url.Values{}
	if realm != "" {
		data.Set("REALM", realm)
	}

	req, err := s.NewRequest("POST", "/api/v1/apps/list", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// Logs retrieves application logs
func (s *SecurePass) Logs(realm, start, end string) (*LogsResponse, error) {
	var obj LogsResponse

	data := url.Values{}
	if realm != "" {
		data.Set("REALM", realm)
	}
	if start != "" {
		data.Set("START", start)
	}
	if end != "" {
		data.Set("END", end)
	}

	req, err := s.NewRequest("POST", "/api/v1/logs/get", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// GroupMember issues requests to /api/v1/groups/member
func (s *SecurePass) GroupMember(user, group string) (*GroupMemberResponse, error) {
	var obj GroupMemberResponse

	data := url.Values{}
	data.Set("USERNAME", user)
	data.Set("GROUP", group)

	req, err := s.NewRequest("POST", "/api/v1/groups/member", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// UserInfo issues requests to /api/v1/users/info
func (s *SecurePass) UserInfo(username string) (*UserInfoResponse, error) {
	var obj UserInfoResponse

	data := url.Values{}
	data.Set("USERNAME", username)

	req, err := s.NewRequest("POST", "/api/v1/users/info", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// UserList issues requests to /api/v1/users/list
func (s *SecurePass) UserList(realm string) (*UserListResponse, error) {
	var obj UserListResponse

	data := url.Values{}
	if realm != "" {
		data.Set("REALM", realm)
	}

	req, err := s.NewRequest("POST", "/api/v1/users/list", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// UserAuth issues requests to /api/v1/users/auth
func (s *SecurePass) UserAuth(username, secret string) (*UserAuthResponse, error) {
	var obj UserAuthResponse

	data := url.Values{}
	data.Set("USERNAME", username)
	data.Set("SECRET", secret)

	req, err := s.NewRequest("POST", "/api/v1/users/auth", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// UserAdd issues requests to /api/v1/users/add
func (s *SecurePass) UserAdd(user *UserDescriptor) (*UserAddResponse, error) {
	var obj UserAddResponse

	data := url.Values{}
	data.Set("USERNAME", user.Username)
	data.Set("NAME", user.Name)
	data.Set("SURNAME", user.Surname)
	data.Set("EMAIL", user.Email)
	data.Set("MOBILE", user.Mobile)
	data.Set("NIN", user.Nin)
	data.Set("RFID", user.Rfid)
	data.Set("MANAGER", user.Manager)

	req, err := s.NewRequest("POST", "/api/v1/users/add", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// UserDel deletes a user from SecurePass
func (s *SecurePass) UserDel(username string) (*Response, error) {
	var obj Response

	data := url.Values{}
	data.Set("USERNAME", username)

	req, err := s.NewRequest("POST", "/api/v1/users/delete", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// Ping reprenets the /api/v1/ping API call
func (s *SecurePass) Ping() (*PingResponse, error) {
	var obj PingResponse

	req, err := s.NewRequest("GET", "/api/v1/ping", nil)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}
<commit_msg>The endpoint returns just an exit code<commit_after>package securepassctl

import (
	"bytes"
	"crypto/tls"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
)

const (
	// DefaultRemote is the default Content-Type header used in HTTP requests
	DefaultRemote = "https://beta.secure-pass.net"
	// ContentType is the default Content-Type header used in HTTP requests
	ContentType = "application/json"
	// UserAgent contains the default User-Agent value used in HTTP requests
	UserAgent = "SecurePass CLI"
)

// DebugLogger collects all debug messages
var DebugLogger = log.New(ioutil.Discard, "", log.LstdFlags)

// SecurePass main object type
type SecurePass struct {
	AppID     string `ini:"APP_ID"`
	AppSecret string `ini:"APP_SECRET"`
	Endpoint  string
}

func (s *SecurePass) setupRequestFieds(req *http.Request) {
	req.Header.Set("Accept", ContentType)
	req.Header.Set("Content-Type", ContentType)
	req.Header.Set("User-Agent", UserAgent)
	req.Header.Set("X-SecurePass-App-ID", s.AppID)
	req.Header.Set("X-SecurePass-App-Secret", s.AppSecret)
}

func (s *SecurePass) makeRequestURL(path string) (string, error) {
	baseURL, _ := url.Parse(s.Endpoint)
	URL, err := url.Parse(path)
	if err != nil {
		return "", err
	}
	return baseURL.ResolveReference(URL).String(), nil
}

// NewRequest initializes and issues an HTTP request to the SecurePass endpoint
func (s *SecurePass) NewRequest(method, path string, data *url.Values) (*http.Request, error) {
	var err error
	var req *http.Request

	URL, err := s.makeRequestURL(path)
	if err != nil {
		return nil, err
	}
	if data != nil {
		req, err = http.NewRequest(method, URL, bytes.NewBufferString(data.Encode()))
	} else {
		req, err = http.NewRequest(method, URL, nil)
	}
	if err != nil {
		return nil, err
	}
	s.setupRequestFieds(req)
	return req, nil
}

// DoRequest issues an HTTP request
func (s *SecurePass) DoRequest(req *http.Request, obj APIResponse, expstatus int) error {
	client := NewClient(nil)
	DebugLogger.Printf("Sending %s request to %s", req.Method, req.URL)
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	if resp.StatusCode != expstatus {
		return fmt.Errorf("%s", resp.Status)
	}
	defer resp.Body.Close()

	err = json.NewDecoder(resp.Body).Decode(obj)
	if err != nil {
		return err
	}
	if obj.ErrorCode() != 0 {
		return fmt.Errorf("%d: %s", obj.ErrorCode(), obj.ErrorMessage())
	}
	return nil
}

// NewClient initialize http.Client with a certain http.Transport
func NewClient(tr *http.Transport) *http.Client {
	// Skip SSL certificate verification
	if tr == nil {
		tr = &http.Transport{TLSClientConfig: &tls.Config{InsecureSkipVerify: true}}
	}
	return &http.Client{Transport: tr}
}

// AppInfo retrieves information on a SecurePass application
func (s *SecurePass) AppInfo(app string) (*AppInfoResponse, error) {
	var obj AppInfoResponse

	data := url.Values{}
	if app != "" {
		data.Set("APP_ID", app)
	}

	req, err := s.NewRequest("POST", "/api/v1/apps/info", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// AppAdd represents /api/v1/apps/add
func (s *SecurePass) AppAdd(app *ApplicationDescriptor) (*AppAddResponse, error) {
	var obj AppAddResponse

	data := url.Values{}
	data.Set("LABEL", app.Label)
	data.Set("WRITE", fmt.Sprintf("%v", app.Write))
	data.Set("PRIVACY", fmt.Sprintf("%v", app.Privacy))
	if app.AllowNetworkIPv4 != "" {
		data.Set("ALLOW_NETWORK_IPv4", app.AllowNetworkIPv4)
	}
	if app.AllowNetworkIPv6 != "" {
		data.Set("ALLOW_NETWORK_IPv6", app.AllowNetworkIPv6)
	}
	if app.Group != "" {
		data.Set("GROUP", app.Group)
	}
	if app.Realm != "" {
		data.Set("REALM", app.Realm)
	}

	req, err := s.NewRequest("POST", "/api/v1/apps/add", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// AppDel deletes an application from SecurePass
func (s *SecurePass) AppDel(app string) (*Response, error) {
	var obj Response

	data := url.Values{}
	if app != "" {
		data.Set("APP_ID", app)
	}

	req, err := s.NewRequest("POST", "/api/v1/apps/delete", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// AppMod represents /api/v1/apps/modify
func (s *SecurePass) AppMod(appID string, app *ApplicationDescriptor) (*Response, error) {
	var obj Response

	data := url.Values{}
	data.Set("APP_ID", appID)
	data.Set("WRITE", fmt.Sprintf("%v", app.Write))
	data.Set("PRIVACY", fmt.Sprintf("%v", app.Privacy))
	if app.Label != "" {
		data.Set("LABEL", app.Label)
	}
	if app.AllowNetworkIPv4 != "" {
		data.Set("ALLOW_NETWORK_IPv4", app.AllowNetworkIPv4)
	}
	if app.AllowNetworkIPv6 != "" {
		data.Set("ALLOW_NETWORK_IPv6", app.AllowNetworkIPv6)
	}
	if app.Group != "" {
		data.Set("GROUP", app.Group)
	}

	req, err := s.NewRequest("POST", "/api/v1/apps/modify", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// AppList retrieves the list of applications available in SecurePass
func (s *SecurePass) AppList(realm string) (*AppListResponse, error) {
	var obj AppListResponse

	data := url.Values{}
	if realm != "" {
		data.Set("REALM", realm)
	}

	req, err := s.NewRequest("POST", "/api/v1/apps/list", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// Logs retrieves application logs
func (s *SecurePass) Logs(realm, start, end string) (*LogsResponse, error) {
	var obj LogsResponse

	data := url.Values{}
	if realm != "" {
		data.Set("REALM", realm)
	}
	if start != "" {
		data.Set("START", start)
	}
	if end != "" {
		data.Set("END", end)
	}

	req, err := s.NewRequest("POST", "/api/v1/logs/get", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// GroupMember issues requests to /api/v1/groups/member
func (s *SecurePass) GroupMember(user, group string) (*GroupMemberResponse, error) {
	var obj GroupMemberResponse

	data := url.Values{}
	data.Set("USERNAME", user)
	data.Set("GROUP", group)

	req, err := s.NewRequest("POST", "/api/v1/groups/member", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// UserInfo issues requests to /api/v1/users/info
func (s *SecurePass) UserInfo(username string) (*UserInfoResponse, error) {
	var obj UserInfoResponse

	data := url.Values{}
	data.Set("USERNAME", username)

	req, err := s.NewRequest("POST", "/api/v1/users/info", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// UserList issues requests to /api/v1/users/list
func (s *SecurePass) UserList(realm string) (*UserListResponse, error) {
	var obj UserListResponse

	data := url.Values{}
	if realm != "" {
		data.Set("REALM", realm)
	}

	req, err := s.NewRequest("POST", "/api/v1/users/list", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// UserAuth issues requests to /api/v1/users/auth
func (s *SecurePass) UserAuth(username, secret string) (*UserAuthResponse, error) {
	var obj UserAuthResponse

	data := url.Values{}
	data.Set("USERNAME", username)
	data.Set("SECRET", secret)

	req, err := s.NewRequest("POST", "/api/v1/users/auth", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// UserAdd issues requests to /api/v1/users/add
func (s *SecurePass) UserAdd(user *UserDescriptor) (*UserAddResponse, error) {
	var obj UserAddResponse

	data := url.Values{}
	data.Set("USERNAME", user.Username)
	data.Set("NAME", user.Name)
	data.Set("SURNAME", user.Surname)
	data.Set("EMAIL", user.Email)
	data.Set("MOBILE", user.Mobile)
	data.Set("NIN", user.Nin)
	data.Set("RFID", user.Rfid)
	data.Set("MANAGER", user.Manager)

	req, err := s.NewRequest("POST", "/api/v1/users/add", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// UserDel deletes a user from SecurePass
func (s *SecurePass) UserDel(username string) (*Response, error) {
	var obj Response

	data := url.Values{}
	data.Set("USERNAME", username)

	req, err := s.NewRequest("POST", "/api/v1/users/delete", &data)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}

// Ping reprenets the /api/v1/ping API call
func (s *SecurePass) Ping() (*PingResponse, error) {
	var obj PingResponse

	req, err := s.NewRequest("GET", "/api/v1/ping", nil)
	if err != nil {
		return nil, err
	}
	err = s.DoRequest(req, &obj, 200)
	return &obj, err
}
<|endoftext|>


<commit_before>package providers

import (
	"errors"
	"fmt"
	"github.com/bitly/go-simplejson"
	"github.com/bitly/oauth2_proxy/api"
	"log"
	"net/http"
	"net/url"
)

type AzureProvider struct {
	*ProviderData
	Tenant string
}

func NewAzureProvider(p *ProviderData) *AzureProvider {
	p.ProviderName = "Azure"

	if p.ProfileURL == nil || p.ProfileURL.String() == "" {
		p.ProfileURL = &url.URL{
			Scheme:   "https",
			Host:     "idoco360train.api.crm9.dynamics.com/",
			RawQuery: "api-version=1.6",

		}
	}
	if p.ProtectedResource == nil || p.ProtectedResource.String() == "" {
		p.ProtectedResource = &url.URL{
			Scheme: "https",
			Host:   "idoco360train.api.crm9.dynamics.com/"
		}
	}
	if p.Scope == "" {
		p.Scope = "openid"
	}

	return &AzureProvider{ProviderData: p}
}

func (p *AzureProvider) Configure(tenant string) {
	p.Tenant = tenant
	if tenant == "" {
		p.Tenant = "idocO360.onmicrosoft.com"
	}

	if p.LoginURL == nil || p.LoginURL.String() == "" {
		p.LoginURL = &url.URL{
			Scheme: "https",
			Host:   "login.windows.net",
			Path:   "/" + p.Tenant + "/oauth2/authorize"}
	}
	if p.RedeemURL == nil || p.RedeemURL.String() == "" {
		p.RedeemURL = &url.URL{
			Scheme: "https",
			Host:   "login.windows.net",
			Path:   "/" + p.Tenant + "/oauth2/token",
		}
	}
}

func getAzureHeader(access_token string) http.Header {
	header := make(http.Header)
	header.Set("Authorization", fmt.Sprintf("Bearer %s", access_token))
	return header
}

func getEmailFromJSON(json *simplejson.Json) (string, error) {
	var email string
	var err error

	email, err = json.Get("mail").String()

	if err != nil || email == "" {
		otherMails, otherMailsErr := json.Get("otherMails").Array()
		if len(otherMails) > 0 {
			email = otherMails[0].(string)
		}
		err = otherMailsErr
	}

	return email, err
}

func (p *AzureProvider) GetEmailAddress(s *SessionState) (string, error) {
	var email string
	var err error

	if s.AccessToken == "" {
		return "", errors.New("missing access token")
	}
	req, err := http.NewRequest("GET", p.ProfileURL.String(), nil)
	if err != nil {
		return "", err
	}
	req.Header = getAzureHeader(s.AccessToken)

	json, err := api.Request(req)

	if err != nil {
		return "", err
	}

	email, err = getEmailFromJSON(json)

	if err == nil && email != "" {
		return email, err
	}

	email, err = json.Get("userPrincipalName").String()

	if err != nil {
		log.Printf("failed making request %s", err)
		return "", err
	}

	if email == "" {
		log.Printf("failed to get email address")
		return "", err
	}

	return email, err
}
<commit_msg>updated2<commit_after>package providers

import (
	"errors"
	"fmt"
	"github.com/bitly/go-simplejson"
	"github.com/bitly/oauth2_proxy/api"
	"log"
	"net/http"
	"net/url"
)

type AzureProvider struct {
	*ProviderData
	Tenant string
}

func NewAzureProvider(p *ProviderData) *AzureProvider {
	p.ProviderName = "Azure"

	if p.ProfileURL == nil || p.ProfileURL.String() == "" {
		p.ProfileURL = &url.URL{
			Scheme:   "https",
			Host:     "idoco360train.api.crm9.dynamics.com/",
			RawQuery: "api-version=1.6",

		}
	}
	if p.ProtectedResource == nil || p.ProtectedResource.String() == "" {
		p.ProtectedResource = &url.URL{
			Scheme: "https",
			Host:   "idoco360train.api.crm9.dynamics.com/",
		}
	}
	if p.Scope == "" {
		p.Scope = "openid"
	}

	return &AzureProvider{ProviderData: p}
}

func (p *AzureProvider) Configure(tenant string) {
	p.Tenant = tenant
	if tenant == "" {
		p.Tenant = "idocO360.onmicrosoft.com"
	}

	if p.LoginURL == nil || p.LoginURL.String() == "" {
		p.LoginURL = &url.URL{
			Scheme: "https",
			Host:   "login.windows.net",
			Path:   "/" + p.Tenant + "/oauth2/authorize"}
	}
	if p.RedeemURL == nil || p.RedeemURL.String() == "" {
		p.RedeemURL = &url.URL{
			Scheme: "https",
			Host:   "login.windows.net",
			Path:   "/" + p.Tenant + "/oauth2/token",
		}
	}
}

func getAzureHeader(access_token string) http.Header {
	header := make(http.Header)
	header.Set("Authorization", fmt.Sprintf("Bearer %s", access_token))
	return header
}

func getEmailFromJSON(json *simplejson.Json) (string, error) {
	var email string
	var err error

	email, err = json.Get("mail").String()

	if err != nil || email == "" {
		otherMails, otherMailsErr := json.Get("otherMails").Array()
		if len(otherMails) > 0 {
			email = otherMails[0].(string)
		}
		err = otherMailsErr
	}

	return email, err
}

func (p *AzureProvider) GetEmailAddress(s *SessionState) (string, error) {
	var email string
	var err error

	if s.AccessToken == "" {
		return "", errors.New("missing access token")
	}
	req, err := http.NewRequest("GET", p.ProfileURL.String(), nil)
	if err != nil {
		return "", err
	}
	req.Header = getAzureHeader(s.AccessToken)

	json, err := api.Request(req)

	if err != nil {
		return "", err
	}

	email, err = getEmailFromJSON(json)

	if err == nil && email != "" {
		return email, err
	}

	email, err = json.Get("userPrincipalName").String()

	if err != nil {
		log.Printf("failed making request %s", err)
		return "", err
	}

	if email == "" {
		log.Printf("failed to get email address")
		return "", err
	}

	return email, err
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"github.com/mingrammer/go-codelab/models"
	"log"
	"net/http"
	"os"
	"sync"
	"strings"
)

const (
	logDir = "log"
	tempLog = "Temp.log"
	accelLog = "Accel.log"
	gyroLog = "Gyro.log"
)

type logContent struct {
	content  string
	location string
}

type TempHandler struct {
	buf chan logContent
}

type GyroHandler struct {
	buf chan logContent
}

type AccelHandler struct {
	buf chan logContent
}

func (m *TempHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	var data models.TempSensor

	decoder := json.NewDecoder(req.Body)
	err := decoder.Decode(&data)
	if err != nil {
		fmt.Println("Somethink wrong")
	}
	defer req.Body.Close()

	m.buf <- logContent{content: fmt.Sprintf("%s", data), location: tempLog}
}

func (m *GyroHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	var data models.GyroSensor

	decoder := json.NewDecoder(req.Body)
	err := decoder.Decode(&data)
	if err != nil {
		fmt.Println("Something wrong")
	}
	defer req.Body.Close()

	m.buf <- logContent{content: fmt.Sprintf("%s", data), location: gyroLog}
}

func (m *AccelHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	var data models.AccelSensor

	decoder := json.NewDecoder(req.Body)
	err := decoder.Decode(&data)
	if err != nil {
		fmt.Println("Something wrong")
	}
	defer req.Body.Close()

	m.buf <- logContent{content: fmt.Sprintf("%s", data), location: accelLog}
}

func fileLogger(m chan logContent) {

	for i := range m {
		joinee := []string{logDir, i.location}
		filePath := strings.Join(joinee, "/")
		fileHandle, err := os.OpenFile(filePath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)

		if err != nil {
			log.Fatal("Error Occured Opening File\n%s", err)
		}

		logger := log.New(fileHandle, "", log.LstdFlags)

		switch i.location {
		case gyroLog:
			logger.Printf("[GyroSensor Data Received]\n%s\n", i.content)
		case accelLog:
			logger.Printf("[AccelSensor Data Received]\n%s\n", i.content)
		case tempLog:
			logger.Printf("[TempSensor Data Received]\n%s\n", i.content)
		}

		defer fileHandle.Close()
	}
}

func main() {
	var wg sync.WaitGroup

	wg.Add(4)

	logBuf := make(chan logContent)
	gyroHander := &GyroHandler{buf: logBuf}
	accelHandler := &AccelHandler{buf: logBuf}
	tempHandler := &TempHandler{buf: logBuf}

	go http.ListenAndServe(":8001", gyroHander)
	go http.ListenAndServe(":8002", accelHandler)
	go http.ListenAndServe(":8003", tempHandler)
	go fileLogger(logBuf)

	wg.Wait()
}
<commit_msg>Apply 'Go fmt'<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"github.com/mingrammer/go-codelab/models"
	"log"
	"net/http"
	"os"
	"strings"
	"sync"
)

const (
	logDir   = "log"
	tempLog  = "Temp.log"
	accelLog = "Accel.log"
	gyroLog  = "Gyro.log"
)

type logContent struct {
	content  string
	location string
}

type TempHandler struct {
	buf chan logContent
}

type GyroHandler struct {
	buf chan logContent
}

type AccelHandler struct {
	buf chan logContent
}

func (m *TempHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	var data models.TempSensor

	decoder := json.NewDecoder(req.Body)
	err := decoder.Decode(&data)
	if err != nil {
		fmt.Println("Somethink wrong")
	}
	defer req.Body.Close()

	m.buf <- logContent{content: fmt.Sprintf("%s", data), location: tempLog}
}

func (m *GyroHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	var data models.GyroSensor

	decoder := json.NewDecoder(req.Body)
	err := decoder.Decode(&data)
	if err != nil {
		fmt.Println("Something wrong")
	}
	defer req.Body.Close()

	m.buf <- logContent{content: fmt.Sprintf("%s", data), location: gyroLog}
}

func (m *AccelHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	var data models.AccelSensor

	decoder := json.NewDecoder(req.Body)
	err := decoder.Decode(&data)
	if err != nil {
		fmt.Println("Something wrong")
	}
	defer req.Body.Close()

	m.buf <- logContent{content: fmt.Sprintf("%s", data), location: accelLog}
}

func fileLogger(m chan logContent) {

	for i := range m {
		joinee := []string{logDir, i.location}
		filePath := strings.Join(joinee, "/")
		fileHandle, err := os.OpenFile(filePath, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)

		if err != nil {
			log.Fatal("Error Occured Opening File\n%s", err)
		}

		logger := log.New(fileHandle, "", log.LstdFlags)

		switch i.location {
		case gyroLog:
			logger.Printf("[GyroSensor Data Received]\n%s\n", i.content)
		case accelLog:
			logger.Printf("[AccelSensor Data Received]\n%s\n", i.content)
		case tempLog:
			logger.Printf("[TempSensor Data Received]\n%s\n", i.content)
		}

		defer fileHandle.Close()
	}
}

func main() {
	var wg sync.WaitGroup

	wg.Add(4)

	logBuf := make(chan logContent)
	gyroHander := &GyroHandler{buf: logBuf}
	accelHandler := &AccelHandler{buf: logBuf}
	tempHandler := &TempHandler{buf: logBuf}

	go http.ListenAndServe(":8001", gyroHander)
	go http.ListenAndServe(":8002", accelHandler)
	go http.ListenAndServe(":8003", tempHandler)
	go fileLogger(logBuf)

	wg.Wait()
}
<|endoftext|>


<commit_before>package paperless

import (
	"bytes"
	"os"
	"reflect"
	"testing"
)

func Test_parseConsts(t *testing.T) {
	type args struct {
		s string
	}
	tests := []struct {
		name string
		args args
		want []string
	}{
		{"No variables", args{"something else"}, []string{}},
		{"Single variable", args{"$a"}, []string{"a"}},
		{"Many variables", args{"$a $b $c"}, []string{"a", "b", "c"}},
		{"Combined", args{"$first$second"}, []string{"first", "second"}},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := parseConsts(tt.args.s); !reflect.DeepEqual(got, tt.want) {
				t.Errorf("parseConsts() = %v, want %v", got, tt.want)
			}
		})
	}
}

func Test_expandConsts(t *testing.T) {
	type args struct {
		s         string
		constants map[string]string
	}
	tests := []struct {
		name string
		args args
		want string
	}{
		{"Empty", args{"", map[string]string{}}, ""},
		{"Single constant", args{"$abc", map[string]string{
			"abc": "something",
		}}, "something"},
		{"Multiple constants", args{"$a$b", map[string]string{
			"a": "some",
			"b": "thing",
		}}, "something"},
		{"Other stuff", args{"$a other $b", map[string]string{
			"a": "some",
			"b": "thing",
		}}, "some other thing"},
		{"Undefined", args{"$a$undefined$b", map[string]string{
			"a": "some",
			"b": "thing",
		}}, "something"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := expandConsts(tt.args.s, tt.args.constants); got != tt.want {
				t.Errorf("expandConsts() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestNewCmdChainScript(t *testing.T) {
	type args struct {
		script string
	}
	tests := []struct {
		name    string
		args    args
		wantC   *CmdChain
		wantErr bool
	}{
		{"Empty", args{""}, &CmdChain{
			Environment: Environment{Constants: map[string]string{}},
		}, false},

		{"Comment and empty line", args{`
# comment`}, &CmdChain{
			Environment: Environment{Constants: map[string]string{}},
		}, false},

		{"Single command", args{"true"}, &CmdChain{
			Links:       []Link{&Cmd{[]string{"true"}}},
			Environment: Environment{Constants: map[string]string{}},
		}, false},

		{"Two commands", args{"true\nfalse"}, &CmdChain{
			Links: []Link{
				&Cmd{[]string{"true"}},
				&Cmd{[]string{"false"}},
			},
			Environment: Environment{Constants: map[string]string{}},
		}, false},

		{"Arguments", args{"true first second"}, &CmdChain{
			Links: []Link{
				&Cmd{[]string{"true", "first", "second"}},
			},
			Environment: Environment{Constants: map[string]string{}},
		}, false},

		{"Quoted arguments", args{"true 'first second'"}, &CmdChain{
			Links: []Link{
				&Cmd{[]string{"true", "first second"}},
			},
			Environment: Environment{Constants: map[string]string{}},
		}, false},

		{"Included a constant", args{"true $variable"}, &CmdChain{
			Environment: Environment{
				Constants: map[string]string{
					"variable": "",
				},
			},
			Links: []Link{
				&Cmd{[]string{"true", "$variable"}},
			},
		}, false},

		{"Command not found", args{"this-command-is-not-found"}, nil, true},

		{"Included a temporary file", args{"true $tmpSomething"}, &CmdChain{
			Environment: Environment{
				Constants: map[string]string{
					"tmpSomething": "",
				},
				TempFiles: []string{"tmpSomething"},
			},
			Links: []Link{
				&Cmd{[]string{"true", "$tmpSomething"}},
			},
		}, false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			gotC, err := NewCmdChainScript(tt.args.script)
			if (err != nil) != tt.wantErr {
				t.Errorf("NewCmdChainScript() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(gotC, tt.wantC) {
				t.Errorf("NewCmdChainScript() = %v, want %v", gotC, tt.wantC)
			}
		})
	}
}

func Test_splitWsQuote(t *testing.T) {
	type args struct {
		s string
	}
	tests := []struct {
		name string
		args args
		want []string
	}{
		{"Empty", args{""}, []string{}},
		{"One item", args{"jep"}, []string{"jep"}},
		{"Two items", args{"jep something"}, []string{"jep", "something"}},
		{"Quoted", args{"'sth abc'"}, []string{"sth abc"}},
		{"Quoted two", args{"'sth' abc"}, []string{"sth", "abc"}},
		{"Mixed quotes", args{"'a b c' \"c  e\""}, []string{"a b c", "c  e"}},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := splitWsQuote(tt.args.s); !reflect.DeepEqual(got, tt.want) {
				t.Errorf("splitWsQuote() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestCmd_Validate(t *testing.T) {
	type fields struct {
		Cmd []string
	}
	type args struct {
		e Environment
	}
	tests := []struct {
		name    string
		fields  fields
		args    args
		wantErr bool
	}{
		{"Proper command", fields{[]string{"true"}},
			args{Environment{
				RootDir: "/",
			}}, false},
		{"Empty command", fields{[]string{""}}, args{}, true},
		// {"LastErr already set", fields{[]string{"true"}}, args{Status{LastErr: errors.New("abc")}}, true},
		{"Command not found", fields{[]string{"command-is-not-found"}}, args{}, true},
		{"Command is allowed", fields{[]string{"true"}},
			args{Environment{
				RootDir: "/",
				AllowedCommands: map[string]bool{
					"true": true,
				},
			}}, false},
		{"Command not allowed", fields{[]string{"true"}},
			args{Environment{
				AllowedCommands: map[string]bool{
					"b": true,
				},
			}}, true},
		{"Constant is defined", fields{[]string{"true", "$something"}},
			args{Environment{
				RootDir: "/",
				Constants: map[string]string{
					"something": "value",
				},
			}}, false},

		{"Constant is not defined", fields{[]string{"true", "$else"}},
			args{Environment{
				RootDir: "/",
			}}, true},

		{"Commands cannot be read from a constant", fields{[]string{"$cmd"}},
			args{Environment{
				Constants: map[string]string{
					"cmd": "true",
				},
			}}, true},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			c := &Cmd{
				Cmd: tt.fields.Cmd,
			}
			if err := c.Validate(&tt.args.e); (err != nil) != tt.wantErr {
				t.Errorf("Cmd.Validate() error = %v, wantErr %v", err, tt.wantErr)
			}
		})
	}
}

func envIsProper(e *Environment) bool {
	return e.initialized && e.validate() == nil
}

func envTempfilesExist(e *Environment) (ret bool) {
	if e.validate() != nil {
		return false
	}

	if len(e.TempFiles) == 0 {
		return false
	}

	ret = true
	for _, n := range e.TempFiles {
		_, ok := e.Constants[n]
		if !ok {
			ret = false
			continue
		}
		_, err := os.Stat(e.Constants[n])
		if err != nil {
			ret = false
		}
	}
	return
}

func TestEnvironment_initEnv(t *testing.T) {
	tests := []struct {
		name     string
		fields   Environment
		wantErr  bool
		validate func(*Environment) bool
	}{
		{"Empty Environment", Environment{}, true, nil},
		{"Already initialized", Environment{
			initialized: true,
		}, false, nil},
		{"Proper but empty", Environment{
			Constants: map[string]string{},
		}, false, envIsProper},
		{"Proper with tempfiles", Environment{
			Constants: map[string]string{
				"a": "",
			},
			TempFiles: []string{
				"a",
			},
		}, false, envTempfilesExist},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var err error
			e := &tt.fields
			if err = e.initEnv(); (err != nil) != tt.wantErr {
				t.Errorf("Environment.initEnv() error = %v, wantErr %v", err, tt.wantErr)
			}

			if tt.validate != nil && !tt.validate(e) {
				t.Errorf("Environment should be proper")
			}

			if err == nil {
				e.deinitEnv()
			}
		})
	}
}

func tempfilesShouldntExist(orig, deinit *Environment) (ret bool) {
	if len(orig.TempFiles) == 0 {
		return false
	}
	for _, n := range orig.TempFiles {
		fname := orig.Constants[n]
		_, err := os.Stat(fname)
		if err == nil {
			return false
		}
	}

	return true
}

func TestEnvironment_deinitEnv(t *testing.T) {
	tests := []struct {
		name       string
		fields     Environment
		shouldInit bool
		wantErr    bool
		validate   func(orig, deinit *Environment) bool
	}{
		{"Empty Environment", Environment{}, false, false, nil},
		{"Proper deinit", Environment{
			Constants: map[string]string{
				"a": "",
			},
			TempFiles: []string{
				"a",
			},
		}, true, false, tempfilesShouldntExist},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			e := &tt.fields
			if tt.shouldInit {
				err := e.initEnv()
				if err != nil {
					t.Errorf("initEnv should succeed")
					return
				}
			}
			orig := *e

			if err := e.deinitEnv(); (err != nil) != tt.wantErr {
				t.Errorf("Environment.deinitEnv() error = %v, wantErr %v", err, tt.wantErr)
			}

			if tt.validate != nil && !tt.validate(&orig, e) {
				t.Errorf("Environment should be deinitialized properly")
			}

		})
	}
}

func TestRunCmdChain(t *testing.T) {
	tests := []struct {
		name       string
		script     string
		consts     map[string]string
		valid      bool
		wantOutput bool
		output     string
		wantErr    bool
	}{
		{"Empty script", "", nil, true, true, "", false},
		{"Echo command", "echo piip", nil, true, true,
			"Running command: echo piip\npiip\n", false},
		{"Echo with a constant", "echo $msg", map[string]string{
			"msg": "piip",
		}, true, true, "Running command: echo piip\npiip\n", false},
		{"Existing temporary file", "echo $tmpmsg\ncat $tmpmsg", nil,
			true, false, "", false},
		{"Failing commands", "true\nfalse\ntrue", nil, true, false, "", true},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {

			ch, err := NewCmdChainScript(tt.script)
			if (err != nil) == tt.valid {
				t.Errorf("NewCmdChainScript() error = %v, valid %v", err, tt.valid)
			}

			s := Status{Environment: ch.Environment}
			if tt.wantOutput {
				s.Log = &bytes.Buffer{}
			}
			if tt.consts != nil {
				s.Constants = tt.consts
			}

			err = RunCmdChain(ch, &s)
			if (err != nil) != tt.wantErr {
				t.Errorf("RunCmdChain() error = %v, wantErr %v", err, tt.wantErr)
			}

			if tt.wantOutput && s.Log.(*bytes.Buffer).String() != tt.output {
				t.Errorf("RunCmdChain() = [%v], want [%v]",
					s.Log.(*bytes.Buffer).String(),
					tt.output)
			}
		})
	}
}
<commit_msg>cmdchain: refactor test to have RootDir defined by default<commit_after>package paperless

import (
	"bytes"
	"os"
	"reflect"
	"testing"
)

func Test_parseConsts(t *testing.T) {
	type args struct {
		s string
	}
	tests := []struct {
		name string
		args args
		want []string
	}{
		{"No variables", args{"something else"}, []string{}},
		{"Single variable", args{"$a"}, []string{"a"}},
		{"Many variables", args{"$a $b $c"}, []string{"a", "b", "c"}},
		{"Combined", args{"$first$second"}, []string{"first", "second"}},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := parseConsts(tt.args.s); !reflect.DeepEqual(got, tt.want) {
				t.Errorf("parseConsts() = %v, want %v", got, tt.want)
			}
		})
	}
}

func Test_expandConsts(t *testing.T) {
	type args struct {
		s         string
		constants map[string]string
	}
	tests := []struct {
		name string
		args args
		want string
	}{
		{"Empty", args{"", map[string]string{}}, ""},
		{"Single constant", args{"$abc", map[string]string{
			"abc": "something",
		}}, "something"},
		{"Multiple constants", args{"$a$b", map[string]string{
			"a": "some",
			"b": "thing",
		}}, "something"},
		{"Other stuff", args{"$a other $b", map[string]string{
			"a": "some",
			"b": "thing",
		}}, "some other thing"},
		{"Undefined", args{"$a$undefined$b", map[string]string{
			"a": "some",
			"b": "thing",
		}}, "something"},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := expandConsts(tt.args.s, tt.args.constants); got != tt.want {
				t.Errorf("expandConsts() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestNewCmdChainScript(t *testing.T) {
	type args struct {
		script string
	}
	tests := []struct {
		name    string
		args    args
		wantC   *CmdChain
		wantErr bool
	}{
		{"Empty", args{""}, &CmdChain{
			Environment: Environment{Constants: map[string]string{}},
		}, false},

		{"Comment and empty line", args{`
# comment`}, &CmdChain{
			Environment: Environment{Constants: map[string]string{}},
		}, false},

		{"Single command", args{"true"}, &CmdChain{
			Links:       []Link{&Cmd{[]string{"true"}}},
			Environment: Environment{Constants: map[string]string{}},
		}, false},

		{"Two commands", args{"true\nfalse"}, &CmdChain{
			Links: []Link{
				&Cmd{[]string{"true"}},
				&Cmd{[]string{"false"}},
			},
			Environment: Environment{Constants: map[string]string{}},
		}, false},

		{"Arguments", args{"true first second"}, &CmdChain{
			Links: []Link{
				&Cmd{[]string{"true", "first", "second"}},
			},
			Environment: Environment{Constants: map[string]string{}},
		}, false},

		{"Quoted arguments", args{"true 'first second'"}, &CmdChain{
			Links: []Link{
				&Cmd{[]string{"true", "first second"}},
			},
			Environment: Environment{Constants: map[string]string{}},
		}, false},

		{"Included a constant", args{"true $variable"}, &CmdChain{
			Environment: Environment{
				Constants: map[string]string{
					"variable": "",
				},
			},
			Links: []Link{
				&Cmd{[]string{"true", "$variable"}},
			},
		}, false},

		{"Command not found", args{"this-command-is-not-found"}, nil, true},

		{"Included a temporary file", args{"true $tmpSomething"}, &CmdChain{
			Environment: Environment{
				Constants: map[string]string{
					"tmpSomething": "",
				},
				TempFiles: []string{"tmpSomething"},
			},
			Links: []Link{
				&Cmd{[]string{"true", "$tmpSomething"}},
			},
		}, false},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			gotC, err := NewCmdChainScript(tt.args.script)
			if (err != nil) != tt.wantErr {
				t.Errorf("NewCmdChainScript() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if !reflect.DeepEqual(gotC, tt.wantC) {
				t.Errorf("NewCmdChainScript() = %v, want %v", gotC, tt.wantC)
			}
		})
	}
}

func Test_splitWsQuote(t *testing.T) {
	type args struct {
		s string
	}
	tests := []struct {
		name string
		args args
		want []string
	}{
		{"Empty", args{""}, []string{}},
		{"One item", args{"jep"}, []string{"jep"}},
		{"Two items", args{"jep something"}, []string{"jep", "something"}},
		{"Quoted", args{"'sth abc'"}, []string{"sth abc"}},
		{"Quoted two", args{"'sth' abc"}, []string{"sth", "abc"}},
		{"Mixed quotes", args{"'a b c' \"c  e\""}, []string{"a b c", "c  e"}},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := splitWsQuote(tt.args.s); !reflect.DeepEqual(got, tt.want) {
				t.Errorf("splitWsQuote() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestCmd_Validate(t *testing.T) {
	type fields struct {
		Cmd []string
	}
	type args struct {
		e Environment
	}
	tests := []struct {
		name    string
		fields  fields
		args    args
		wantErr bool
	}{
		{"Proper command", fields{[]string{"true"}}, args{}, false},
		{"Empty command", fields{[]string{""}}, args{}, true},
		// {"LastErr already set", fields{[]string{"true"}}, args{Status{LastErr: errors.New("abc")}}, true},
		{"Command not found", fields{[]string{"command-is-not-found"}}, args{}, true},
		{"Command is allowed", fields{[]string{"true"}},
			args{Environment{
				AllowedCommands: map[string]bool{
					"true": true,
				},
			}}, false},
		{"Command not allowed", fields{[]string{"true"}},
			args{Environment{
				AllowedCommands: map[string]bool{
					"b": true,
				},
			}}, true},
		{"Constant is defined", fields{[]string{"true", "$something"}},
			args{Environment{
				Constants: map[string]string{
					"something": "value",
				},
			}}, false},

		{"Constant is not defined", fields{[]string{"true", "$else"}}, args{}, true},

		{"Commands cannot be read from a constant", fields{[]string{"$cmd"}},
			args{Environment{
				Constants: map[string]string{
					"cmd": "true",
				},
			}}, true},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			c := &Cmd{
				Cmd: tt.fields.Cmd,
			}
			tt.args.e.RootDir = "/"

			if err := c.Validate(&tt.args.e); (err != nil) != tt.wantErr {
				t.Errorf("Cmd.Validate() error = %v, wantErr %v", err, tt.wantErr)
			}
		})
	}
}

func envIsProper(e *Environment) bool {
	return e.initialized && e.validate() == nil
}

func envTempfilesExist(e *Environment) (ret bool) {
	if e.validate() != nil {
		return false
	}

	if len(e.TempFiles) == 0 {
		return false
	}

	ret = true
	for _, n := range e.TempFiles {
		_, ok := e.Constants[n]
		if !ok {
			ret = false
			continue
		}
		_, err := os.Stat(e.Constants[n])
		if err != nil {
			ret = false
		}
	}
	return
}

func TestEnvironment_initEnv(t *testing.T) {
	tests := []struct {
		name     string
		fields   Environment
		wantErr  bool
		validate func(*Environment) bool
	}{
		{"Empty Environment", Environment{}, true, nil},
		{"Already initialized", Environment{
			initialized: true,
		}, false, nil},
		{"Proper but empty", Environment{
			Constants: map[string]string{},
		}, false, envIsProper},
		{"Proper with tempfiles", Environment{
			Constants: map[string]string{
				"a": "",
			},
			TempFiles: []string{
				"a",
			},
		}, false, envTempfilesExist},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var err error
			e := &tt.fields
			if err = e.initEnv(); (err != nil) != tt.wantErr {
				t.Errorf("Environment.initEnv() error = %v, wantErr %v", err, tt.wantErr)
			}

			if tt.validate != nil && !tt.validate(e) {
				t.Errorf("Environment should be proper")
			}

			if err == nil {
				e.deinitEnv()
			}
		})
	}
}

func tempfilesShouldntExist(orig, deinit *Environment) (ret bool) {
	if len(orig.TempFiles) == 0 {
		return false
	}
	for _, n := range orig.TempFiles {
		fname := orig.Constants[n]
		_, err := os.Stat(fname)
		if err == nil {
			return false
		}
	}

	return true
}

func TestEnvironment_deinitEnv(t *testing.T) {
	tests := []struct {
		name       string
		fields     Environment
		shouldInit bool
		wantErr    bool
		validate   func(orig, deinit *Environment) bool
	}{
		{"Empty Environment", Environment{}, false, false, nil},
		{"Proper deinit", Environment{
			Constants: map[string]string{
				"a": "",
			},
			TempFiles: []string{
				"a",
			},
		}, true, false, tempfilesShouldntExist},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			e := &tt.fields
			if tt.shouldInit {
				err := e.initEnv()
				if err != nil {
					t.Errorf("initEnv should succeed")
					return
				}
			}
			orig := *e

			if err := e.deinitEnv(); (err != nil) != tt.wantErr {
				t.Errorf("Environment.deinitEnv() error = %v, wantErr %v", err, tt.wantErr)
			}

			if tt.validate != nil && !tt.validate(&orig, e) {
				t.Errorf("Environment should be deinitialized properly")
			}

		})
	}
}

func TestRunCmdChain(t *testing.T) {
	tests := []struct {
		name       string
		script     string
		consts     map[string]string
		valid      bool
		wantOutput bool
		output     string
		wantErr    bool
	}{
		{"Empty script", "", nil, true, true, "", false},
		{"Echo command", "echo piip", nil, true, true,
			"Running command: echo piip\npiip\n", false},
		{"Echo with a constant", "echo $msg", map[string]string{
			"msg": "piip",
		}, true, true, "Running command: echo piip\npiip\n", false},
		{"Existing temporary file", "echo $tmpmsg\ncat $tmpmsg", nil,
			true, false, "", false},
		{"Failing commands", "true\nfalse\ntrue", nil, true, false, "", true},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {

			ch, err := NewCmdChainScript(tt.script)
			if (err != nil) == tt.valid {
				t.Errorf("NewCmdChainScript() error = %v, valid %v", err, tt.valid)
			}

			s := Status{Environment: ch.Environment}
			if tt.wantOutput {
				s.Log = &bytes.Buffer{}
			}
			if tt.consts != nil {
				s.Constants = tt.consts
			}

			err = RunCmdChain(ch, &s)
			if (err != nil) != tt.wantErr {
				t.Errorf("RunCmdChain() error = %v, wantErr %v", err, tt.wantErr)
			}

			if tt.wantOutput && s.Log.(*bytes.Buffer).String() != tt.output {
				t.Errorf("RunCmdChain() = [%v], want [%v]",
					s.Log.(*bytes.Buffer).String(),
					tt.output)
			}
		})
	}
}
<|endoftext|>


<commit_before>package common

import (
	"archive/tar"
	"compress/gzip"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/appc/docker2aci/lib/types"
	"github.com/appc/docker2aci/lib/util"
	"github.com/appc/docker2aci/tarball"
	"github.com/appc/spec/aci"
	"github.com/appc/spec/schema"
	appctypes "github.com/appc/spec/schema/types"
)

const (
	defaultTag    = "latest"
	schemaVersion = "0.5.1"
)

func ParseDockerURL(arg string) *types.ParsedDockerURL {
	if arg == "" {
		return nil
	}

	taglessRemote, tag := parseRepositoryTag(arg)
	if tag == "" {
		tag = defaultTag
	}
	indexURL, imageName := SplitReposName(taglessRemote)

	return &types.ParsedDockerURL{
		IndexURL:  indexURL,
		ImageName: imageName,
		Tag:       tag,
	}
}

func GenerateACI(layerData types.DockerImageData, dockerURL *types.ParsedDockerURL, outputDir string, layerFile *os.File, curPwl []string, compress bool) (string, *schema.ImageManifest, error) {
	manifest, err := GenerateManifest(layerData, dockerURL)
	if err != nil {
		return "", nil, fmt.Errorf("error generating the manifest: %v", err)
	}

	imageName := strings.Replace(dockerURL.ImageName, "/", "-", -1)
	aciPath := imageName + "-" + layerData.ID
	if dockerURL.Tag != "" {
		aciPath += "-" + dockerURL.Tag
	}
	if layerData.OS != "" {
		aciPath += "-" + layerData.OS
		if layerData.Architecture != "" {
			aciPath += "-" + layerData.Architecture
		}
	}
	aciPath += ".aci"

	aciPath = path.Join(outputDir, aciPath)
	manifest, err = writeACI(layerFile, *manifest, curPwl, aciPath, compress)
	if err != nil {
		return "", nil, fmt.Errorf("error writing ACI: %v", err)
	}

	if err := ValidateACI(aciPath); err != nil {
		return "", nil, fmt.Errorf("invalid ACI generated: %v", err)
	}

	return aciPath, manifest, nil
}

func ValidateACI(aciPath string) error {
	aciFile, err := os.Open(aciPath)
	if err != nil {
		return err
	}
	defer aciFile.Close()

	reader, err := aci.NewCompressedTarReader(aciFile)
	if err != nil {
		return err
	}

	if err := aci.ValidateArchive(reader); err != nil {
		return err
	}

	return nil
}

func GenerateManifest(layerData types.DockerImageData, dockerURL *types.ParsedDockerURL) (*schema.ImageManifest, error) {
	dockerConfig := layerData.Config
	genManifest := &schema.ImageManifest{}

	appURL := ""
	// omit docker hub index URL in app name
	if dockerURL.IndexURL != defaultIndex {
		appURL = dockerURL.IndexURL + "/"
	}
	appURL += dockerURL.ImageName + "-" + layerData.ID
	appURL, err := appctypes.SanitizeACIdentifier(appURL)
	if err != nil {
		return nil, err
	}
	name := appctypes.MustACIdentifier(appURL)
	genManifest.Name = *name

	acVersion, err := appctypes.NewSemVer(schemaVersion)
	if err != nil {
		panic("invalid appc spec version")
	}
	genManifest.ACVersion = *acVersion

	genManifest.ACKind = appctypes.ACKind(schema.ImageManifestKind)

	var (
		labels       appctypes.Labels
		parentLabels appctypes.Labels
		annotations  appctypes.Annotations
	)

	layer := appctypes.MustACIdentifier("layer")
	labels = append(labels, appctypes.Label{Name: *layer, Value: layerData.ID})

	tag := dockerURL.Tag
	version := appctypes.MustACIdentifier("version")
	labels = append(labels, appctypes.Label{Name: *version, Value: tag})

	if layerData.OS != "" {
		os := appctypes.MustACIdentifier("os")
		labels = append(labels, appctypes.Label{Name: *os, Value: layerData.OS})
		parentLabels = append(parentLabels, appctypes.Label{Name: *os, Value: layerData.OS})

		if layerData.Architecture != "" {
			arch := appctypes.MustACIdentifier("arch")
			parentLabels = append(parentLabels, appctypes.Label{Name: *arch, Value: layerData.Architecture})
		}
	}

	if layerData.Author != "" {
		authorsKey := appctypes.MustACName("authors")
		annotations = append(annotations, appctypes.Annotation{Name: *authorsKey, Value: layerData.Author})
	}
	epoch := time.Unix(0, 0)
	if !layerData.Created.Equal(epoch) {
		createdKey := appctypes.MustACName("created")
		annotations = append(annotations, appctypes.Annotation{Name: *createdKey, Value: layerData.Created.Format(time.RFC3339)})
	}
	if layerData.Comment != "" {
		commentKey := appctypes.MustACName("docker-comment")
		annotations = append(annotations, appctypes.Annotation{Name: *commentKey, Value: layerData.Comment})
	}

	genManifest.Labels = labels
	genManifest.Annotations = annotations

	if dockerConfig != nil {
		exec := getExecCommand(dockerConfig.Entrypoint, dockerConfig.Cmd)
		if exec != nil {
			user, group := parseDockerUser(dockerConfig.User)
			var env appctypes.Environment
			for _, v := range dockerConfig.Env {
				parts := strings.SplitN(v, "=", 2)
				env.Set(parts[0], parts[1])
			}
			app := &appctypes.App{
				Exec:             exec,
				User:             user,
				Group:            group,
				Environment:      env,
				WorkingDirectory: dockerConfig.WorkingDir,
			}

			app.MountPoints, err = convertVolumesToMPs(dockerConfig.Volumes)
			if err != nil {
				return nil, err
			}

			app.Ports, err = convertPorts(dockerConfig.ExposedPorts, dockerConfig.PortSpecs)
			if err != nil {
				return nil, err
			}

			genManifest.App = app
		}
	}

	if layerData.Parent != "" {
		parentImageNameString := dockerURL.IndexURL + "/" + dockerURL.ImageName + "-" + layerData.Parent
		parentImageNameString, err := appctypes.SanitizeACIdentifier(parentImageNameString)
		if err != nil {
			return nil, err
		}
		parentImageName := appctypes.MustACIdentifier(parentImageNameString)

		genManifest.Dependencies = append(genManifest.Dependencies, appctypes.Dependency{ImageName: *parentImageName, Labels: parentLabels})
	}

	return genManifest, nil
}

func convertPorts(dockerExposedPorts map[string]struct{}, dockerPortSpecs []string) ([]appctypes.Port, error) {
	ports := []appctypes.Port{}

	for ep := range dockerExposedPorts {
		appcPort, err := parseDockerPort(ep)
		if err != nil {
			return nil, err
		}
		ports = append(ports, *appcPort)
	}

	if dockerExposedPorts == nil && dockerPortSpecs != nil {
		util.Debug("warning: docker image uses deprecated PortSpecs field")
		for _, ep := range dockerPortSpecs {
			appcPort, err := parseDockerPort(ep)
			if err != nil {
				return nil, err
			}
			ports = append(ports, *appcPort)
		}
	}

	return ports, nil
}

func parseDockerPort(dockerPort string) (*appctypes.Port, error) {
	var portString string
	proto := "tcp"
	sp := strings.Split(dockerPort, "/")
	if len(sp) < 2 {
		portString = dockerPort
	} else {
		proto = sp[1]
		portString = sp[0]
	}

	port, err := strconv.ParseUint(portString, 10, 0)
	if err != nil {
		return nil, fmt.Errorf("error parsing port %q: %v", portString, err)
	}

	sn, err := appctypes.SanitizeACName(dockerPort)
	if err != nil {
		return nil, err
	}

	appcPort := &appctypes.Port{
		Name:     *appctypes.MustACName(sn),
		Protocol: proto,
		Port:     uint(port),
	}

	return appcPort, nil
}

func convertVolumesToMPs(dockerVolumes map[string]struct{}) ([]appctypes.MountPoint, error) {
	mps := []appctypes.MountPoint{}
	dup := make(map[string]int)

	for p := range dockerVolumes {
		n := filepath.Join("volume", p)
		sn, err := appctypes.SanitizeACName(n)
		if err != nil {
			return nil, err
		}

		// check for duplicate names
		if i, ok := dup[sn]; ok {
			dup[sn] = i + 1
			sn = fmt.Sprintf("%s-%d", sn, i)
		} else {
			dup[sn] = 1
		}

		mp := appctypes.MountPoint{
			Name: *appctypes.MustACName(sn),
			Path: p,
		}

		mps = append(mps, mp)
	}

	return mps, nil
}

func writeACI(layer io.ReadSeeker, manifest schema.ImageManifest, curPwl []string, output string, compress bool) (*schema.ImageManifest, error) {
	aciFile, err := os.Create(output)
	if err != nil {
		return nil, fmt.Errorf("error creating ACI file: %v", err)
	}
	defer aciFile.Close()

	var w io.WriteCloser = aciFile
	if compress {
		w = gzip.NewWriter(aciFile)
		defer w.Close()
	}
	trw := tar.NewWriter(w)
	defer trw.Close()

	if err := WriteRootfsDir(trw); err != nil {
		return nil, fmt.Errorf("error writing rootfs entry: %v", err)
	}

	var whiteouts []string
	convWalker := func(t *tarball.TarFile) error {
		name := t.Name()
		if name == "./" {
			return nil
		}
		t.Header.Name = path.Join("rootfs", name)
		absolutePath := strings.TrimPrefix(t.Header.Name, "rootfs")
		if strings.Contains(t.Header.Name, "/.wh.") {
			whiteouts = append(whiteouts, strings.Replace(absolutePath, ".wh.", "", 1))
			return nil
		}
		if t.Header.Typeflag == tar.TypeLink {
			t.Header.Linkname = path.Join("rootfs", t.Linkname())
		}

		if err := trw.WriteHeader(t.Header); err != nil {
			return err
		}
		if _, err := io.Copy(trw, t.TarStream); err != nil {
			return err
		}

		if !util.In(curPwl, absolutePath) {
			curPwl = append(curPwl, absolutePath)
		}

		return nil
	}
	reader, err := aci.NewCompressedTarReader(layer)
	if err == nil {
		// write files in rootfs/
		if err := tarball.Walk(*reader, convWalker); err != nil {
			return nil, err
		}
	} else {
		// ignore errors: empty layers in tars generated by docker save are not
		// valid tar files so we ignore errors trying to open them. Converted
		// ACIs will have the manifest and an empty rootfs directory in any
		// case.
	}
	newPwl := subtractWhiteouts(curPwl, whiteouts)

	manifest.PathWhitelist = newPwl
	if err := WriteManifest(trw, manifest); err != nil {
		return nil, fmt.Errorf("error writing manifest: %v", err)
	}

	return &manifest, nil
}

func getExecCommand(entrypoint []string, cmd []string) appctypes.Exec {
	var command []string
	if entrypoint == nil && cmd == nil {
		return nil
	}
	command = append(entrypoint, cmd...)
	// non-absolute paths are not allowed, fallback to "/bin/sh -c command"
	if len(command) > 0 && !filepath.IsAbs(command[0]) {
		command_prefix := []string{"/bin/sh", "-c"}
		quoted_command := util.Quote(command)
		command = append(command_prefix, strings.Join(quoted_command, " "))
	}
	return command
}

func parseDockerUser(dockerUser string) (string, string) {
	// if the docker user is empty assume root user and group
	if dockerUser == "" {
		return "0", "0"
	}

	dockerUserParts := strings.Split(dockerUser, ":")

	// when only the user is given, the docker spec says that the default and
	// supplementary groups of the user in /etc/passwd should be applied.
	// Assume root group for now in this case.
	if len(dockerUserParts) < 2 {
		return dockerUserParts[0], "0"
	}

	return dockerUserParts[0], dockerUserParts[1]
}

func subtractWhiteouts(pathWhitelist []string, whiteouts []string) []string {
	matchPaths := []string{}
	for _, path := range pathWhitelist {
		// If one of the parent dirs of the current path matches the
		// whiteout then also this path should be removed
		curPath := path
		for curPath != "/" {
			for _, whiteout := range whiteouts {
				if curPath == whiteout {
					matchPaths = append(matchPaths, path)
				}
			}
			curPath = filepath.Dir(curPath)
		}
	}
	for _, matchPath := range matchPaths {
		idx := util.IndexOf(pathWhitelist, matchPath)
		if idx != -1 {
			pathWhitelist = append(pathWhitelist[:idx], pathWhitelist[idx+1:]...)
		}
	}

	return pathWhitelist
}

func WriteManifest(outputWriter *tar.Writer, manifest schema.ImageManifest) error {
	b, err := json.Marshal(manifest)
	if err != nil {
		return err
	}

	hdr := getGenericTarHeader()
	hdr.Name = "manifest"
	hdr.Mode = 0644
	hdr.Size = int64(len(b))
	hdr.Typeflag = tar.TypeReg

	if err := outputWriter.WriteHeader(hdr); err != nil {
		return err
	}
	if _, err := outputWriter.Write(b); err != nil {
		return err
	}

	return nil
}

func WriteRootfsDir(tarWriter *tar.Writer) error {
	hdr := getGenericTarHeader()
	hdr.Name = "rootfs"
	hdr.Mode = 0755
	hdr.Size = int64(0)
	hdr.Typeflag = tar.TypeDir

	return tarWriter.WriteHeader(hdr)
}

func getGenericTarHeader() *tar.Header {
	// FIXME(iaguis) Use docker image time instead of the Unix Epoch?
	hdr := &tar.Header{
		Uid:        0,
		Gid:        0,
		ModTime:    time.Unix(0, 0),
		Uname:      "0",
		Gname:      "0",
		ChangeTime: time.Unix(0, 0),
	}

	return hdr
}
<commit_msg>common: fix build by changing annotation's type to ACIdentifier<commit_after>package common

import (
	"archive/tar"
	"compress/gzip"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/appc/docker2aci/lib/types"
	"github.com/appc/docker2aci/lib/util"
	"github.com/appc/docker2aci/tarball"
	"github.com/appc/spec/aci"
	"github.com/appc/spec/schema"
	appctypes "github.com/appc/spec/schema/types"
)

const (
	defaultTag    = "latest"
	schemaVersion = "0.5.1"
)

func ParseDockerURL(arg string) *types.ParsedDockerURL {
	if arg == "" {
		return nil
	}

	taglessRemote, tag := parseRepositoryTag(arg)
	if tag == "" {
		tag = defaultTag
	}
	indexURL, imageName := SplitReposName(taglessRemote)

	return &types.ParsedDockerURL{
		IndexURL:  indexURL,
		ImageName: imageName,
		Tag:       tag,
	}
}

func GenerateACI(layerData types.DockerImageData, dockerURL *types.ParsedDockerURL, outputDir string, layerFile *os.File, curPwl []string, compress bool) (string, *schema.ImageManifest, error) {
	manifest, err := GenerateManifest(layerData, dockerURL)
	if err != nil {
		return "", nil, fmt.Errorf("error generating the manifest: %v", err)
	}

	imageName := strings.Replace(dockerURL.ImageName, "/", "-", -1)
	aciPath := imageName + "-" + layerData.ID
	if dockerURL.Tag != "" {
		aciPath += "-" + dockerURL.Tag
	}
	if layerData.OS != "" {
		aciPath += "-" + layerData.OS
		if layerData.Architecture != "" {
			aciPath += "-" + layerData.Architecture
		}
	}
	aciPath += ".aci"

	aciPath = path.Join(outputDir, aciPath)
	manifest, err = writeACI(layerFile, *manifest, curPwl, aciPath, compress)
	if err != nil {
		return "", nil, fmt.Errorf("error writing ACI: %v", err)
	}

	if err := ValidateACI(aciPath); err != nil {
		return "", nil, fmt.Errorf("invalid ACI generated: %v", err)
	}

	return aciPath, manifest, nil
}

func ValidateACI(aciPath string) error {
	aciFile, err := os.Open(aciPath)
	if err != nil {
		return err
	}
	defer aciFile.Close()

	reader, err := aci.NewCompressedTarReader(aciFile)
	if err != nil {
		return err
	}

	if err := aci.ValidateArchive(reader); err != nil {
		return err
	}

	return nil
}

func GenerateManifest(layerData types.DockerImageData, dockerURL *types.ParsedDockerURL) (*schema.ImageManifest, error) {
	dockerConfig := layerData.Config
	genManifest := &schema.ImageManifest{}

	appURL := ""
	// omit docker hub index URL in app name
	if dockerURL.IndexURL != defaultIndex {
		appURL = dockerURL.IndexURL + "/"
	}
	appURL += dockerURL.ImageName + "-" + layerData.ID
	appURL, err := appctypes.SanitizeACIdentifier(appURL)
	if err != nil {
		return nil, err
	}
	name := appctypes.MustACIdentifier(appURL)
	genManifest.Name = *name

	acVersion, err := appctypes.NewSemVer(schemaVersion)
	if err != nil {
		panic("invalid appc spec version")
	}
	genManifest.ACVersion = *acVersion

	genManifest.ACKind = appctypes.ACKind(schema.ImageManifestKind)

	var (
		labels       appctypes.Labels
		parentLabels appctypes.Labels
		annotations  appctypes.Annotations
	)

	layer := appctypes.MustACIdentifier("layer")
	labels = append(labels, appctypes.Label{Name: *layer, Value: layerData.ID})

	tag := dockerURL.Tag
	version := appctypes.MustACIdentifier("version")
	labels = append(labels, appctypes.Label{Name: *version, Value: tag})

	if layerData.OS != "" {
		os := appctypes.MustACIdentifier("os")
		labels = append(labels, appctypes.Label{Name: *os, Value: layerData.OS})
		parentLabels = append(parentLabels, appctypes.Label{Name: *os, Value: layerData.OS})

		if layerData.Architecture != "" {
			arch := appctypes.MustACIdentifier("arch")
			parentLabels = append(parentLabels, appctypes.Label{Name: *arch, Value: layerData.Architecture})
		}
	}

	if layerData.Author != "" {
		authorsKey := appctypes.MustACIdentifier("authors")
		annotations = append(annotations, appctypes.Annotation{Name: *authorsKey, Value: layerData.Author})
	}
	epoch := time.Unix(0, 0)
	if !layerData.Created.Equal(epoch) {
		createdKey := appctypes.MustACIdentifier("created")
		annotations = append(annotations, appctypes.Annotation{Name: *createdKey, Value: layerData.Created.Format(time.RFC3339)})
	}
	if layerData.Comment != "" {
		commentKey := appctypes.MustACIdentifier("docker-comment")
		annotations = append(annotations, appctypes.Annotation{Name: *commentKey, Value: layerData.Comment})
	}

	genManifest.Labels = labels
	genManifest.Annotations = annotations

	if dockerConfig != nil {
		exec := getExecCommand(dockerConfig.Entrypoint, dockerConfig.Cmd)
		if exec != nil {
			user, group := parseDockerUser(dockerConfig.User)
			var env appctypes.Environment
			for _, v := range dockerConfig.Env {
				parts := strings.SplitN(v, "=", 2)
				env.Set(parts[0], parts[1])
			}
			app := &appctypes.App{
				Exec:             exec,
				User:             user,
				Group:            group,
				Environment:      env,
				WorkingDirectory: dockerConfig.WorkingDir,
			}

			app.MountPoints, err = convertVolumesToMPs(dockerConfig.Volumes)
			if err != nil {
				return nil, err
			}

			app.Ports, err = convertPorts(dockerConfig.ExposedPorts, dockerConfig.PortSpecs)
			if err != nil {
				return nil, err
			}

			genManifest.App = app
		}
	}

	if layerData.Parent != "" {
		parentImageNameString := dockerURL.IndexURL + "/" + dockerURL.ImageName + "-" + layerData.Parent
		parentImageNameString, err := appctypes.SanitizeACIdentifier(parentImageNameString)
		if err != nil {
			return nil, err
		}
		parentImageName := appctypes.MustACIdentifier(parentImageNameString)

		genManifest.Dependencies = append(genManifest.Dependencies, appctypes.Dependency{ImageName: *parentImageName, Labels: parentLabels})
	}

	return genManifest, nil
}

func convertPorts(dockerExposedPorts map[string]struct{}, dockerPortSpecs []string) ([]appctypes.Port, error) {
	ports := []appctypes.Port{}

	for ep := range dockerExposedPorts {
		appcPort, err := parseDockerPort(ep)
		if err != nil {
			return nil, err
		}
		ports = append(ports, *appcPort)
	}

	if dockerExposedPorts == nil && dockerPortSpecs != nil {
		util.Debug("warning: docker image uses deprecated PortSpecs field")
		for _, ep := range dockerPortSpecs {
			appcPort, err := parseDockerPort(ep)
			if err != nil {
				return nil, err
			}
			ports = append(ports, *appcPort)
		}
	}

	return ports, nil
}

func parseDockerPort(dockerPort string) (*appctypes.Port, error) {
	var portString string
	proto := "tcp"
	sp := strings.Split(dockerPort, "/")
	if len(sp) < 2 {
		portString = dockerPort
	} else {
		proto = sp[1]
		portString = sp[0]
	}

	port, err := strconv.ParseUint(portString, 10, 0)
	if err != nil {
		return nil, fmt.Errorf("error parsing port %q: %v", portString, err)
	}

	sn, err := appctypes.SanitizeACName(dockerPort)
	if err != nil {
		return nil, err
	}

	appcPort := &appctypes.Port{
		Name:     *appctypes.MustACName(sn),
		Protocol: proto,
		Port:     uint(port),
	}

	return appcPort, nil
}

func convertVolumesToMPs(dockerVolumes map[string]struct{}) ([]appctypes.MountPoint, error) {
	mps := []appctypes.MountPoint{}
	dup := make(map[string]int)

	for p := range dockerVolumes {
		n := filepath.Join("volume", p)
		sn, err := appctypes.SanitizeACName(n)
		if err != nil {
			return nil, err
		}

		// check for duplicate names
		if i, ok := dup[sn]; ok {
			dup[sn] = i + 1
			sn = fmt.Sprintf("%s-%d", sn, i)
		} else {
			dup[sn] = 1
		}

		mp := appctypes.MountPoint{
			Name: *appctypes.MustACName(sn),
			Path: p,
		}

		mps = append(mps, mp)
	}

	return mps, nil
}

func writeACI(layer io.ReadSeeker, manifest schema.ImageManifest, curPwl []string, output string, compress bool) (*schema.ImageManifest, error) {
	aciFile, err := os.Create(output)
	if err != nil {
		return nil, fmt.Errorf("error creating ACI file: %v", err)
	}
	defer aciFile.Close()

	var w io.WriteCloser = aciFile
	if compress {
		w = gzip.NewWriter(aciFile)
		defer w.Close()
	}
	trw := tar.NewWriter(w)
	defer trw.Close()

	if err := WriteRootfsDir(trw); err != nil {
		return nil, fmt.Errorf("error writing rootfs entry: %v", err)
	}

	var whiteouts []string
	convWalker := func(t *tarball.TarFile) error {
		name := t.Name()
		if name == "./" {
			return nil
		}
		t.Header.Name = path.Join("rootfs", name)
		absolutePath := strings.TrimPrefix(t.Header.Name, "rootfs")
		if strings.Contains(t.Header.Name, "/.wh.") {
			whiteouts = append(whiteouts, strings.Replace(absolutePath, ".wh.", "", 1))
			return nil
		}
		if t.Header.Typeflag == tar.TypeLink {
			t.Header.Linkname = path.Join("rootfs", t.Linkname())
		}

		if err := trw.WriteHeader(t.Header); err != nil {
			return err
		}
		if _, err := io.Copy(trw, t.TarStream); err != nil {
			return err
		}

		if !util.In(curPwl, absolutePath) {
			curPwl = append(curPwl, absolutePath)
		}

		return nil
	}
	reader, err := aci.NewCompressedTarReader(layer)
	if err == nil {
		// write files in rootfs/
		if err := tarball.Walk(*reader, convWalker); err != nil {
			return nil, err
		}
	} else {
		// ignore errors: empty layers in tars generated by docker save are not
		// valid tar files so we ignore errors trying to open them. Converted
		// ACIs will have the manifest and an empty rootfs directory in any
		// case.
	}
	newPwl := subtractWhiteouts(curPwl, whiteouts)

	manifest.PathWhitelist = newPwl
	if err := WriteManifest(trw, manifest); err != nil {
		return nil, fmt.Errorf("error writing manifest: %v", err)
	}

	return &manifest, nil
}

func getExecCommand(entrypoint []string, cmd []string) appctypes.Exec {
	var command []string
	if entrypoint == nil && cmd == nil {
		return nil
	}
	command = append(entrypoint, cmd...)
	// non-absolute paths are not allowed, fallback to "/bin/sh -c command"
	if len(command) > 0 && !filepath.IsAbs(command[0]) {
		command_prefix := []string{"/bin/sh", "-c"}
		quoted_command := util.Quote(command)
		command = append(command_prefix, strings.Join(quoted_command, " "))
	}
	return command
}

func parseDockerUser(dockerUser string) (string, string) {
	// if the docker user is empty assume root user and group
	if dockerUser == "" {
		return "0", "0"
	}

	dockerUserParts := strings.Split(dockerUser, ":")

	// when only the user is given, the docker spec says that the default and
	// supplementary groups of the user in /etc/passwd should be applied.
	// Assume root group for now in this case.
	if len(dockerUserParts) < 2 {
		return dockerUserParts[0], "0"
	}

	return dockerUserParts[0], dockerUserParts[1]
}

func subtractWhiteouts(pathWhitelist []string, whiteouts []string) []string {
	matchPaths := []string{}
	for _, path := range pathWhitelist {
		// If one of the parent dirs of the current path matches the
		// whiteout then also this path should be removed
		curPath := path
		for curPath != "/" {
			for _, whiteout := range whiteouts {
				if curPath == whiteout {
					matchPaths = append(matchPaths, path)
				}
			}
			curPath = filepath.Dir(curPath)
		}
	}
	for _, matchPath := range matchPaths {
		idx := util.IndexOf(pathWhitelist, matchPath)
		if idx != -1 {
			pathWhitelist = append(pathWhitelist[:idx], pathWhitelist[idx+1:]...)
		}
	}

	return pathWhitelist
}

func WriteManifest(outputWriter *tar.Writer, manifest schema.ImageManifest) error {
	b, err := json.Marshal(manifest)
	if err != nil {
		return err
	}

	hdr := getGenericTarHeader()
	hdr.Name = "manifest"
	hdr.Mode = 0644
	hdr.Size = int64(len(b))
	hdr.Typeflag = tar.TypeReg

	if err := outputWriter.WriteHeader(hdr); err != nil {
		return err
	}
	if _, err := outputWriter.Write(b); err != nil {
		return err
	}

	return nil
}

func WriteRootfsDir(tarWriter *tar.Writer) error {
	hdr := getGenericTarHeader()
	hdr.Name = "rootfs"
	hdr.Mode = 0755
	hdr.Size = int64(0)
	hdr.Typeflag = tar.TypeDir

	return tarWriter.WriteHeader(hdr)
}

func getGenericTarHeader() *tar.Header {
	// FIXME(iaguis) Use docker image time instead of the Unix Epoch?
	hdr := &tar.Header{
		Uid:        0,
		Gid:        0,
		ModTime:    time.Unix(0, 0),
		Uname:      "0",
		Gname:      "0",
		ChangeTime: time.Unix(0, 0),
	}

	return hdr
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Gravitational, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package native

import (
	"testing"
	"time"

	"golang.org/x/crypto/ssh"

	"github.com/gravitational/teleport"
	"github.com/gravitational/teleport/lib/auth/test"
	"github.com/gravitational/teleport/lib/services"
	"github.com/gravitational/teleport/lib/utils"

	. "gopkg.in/check.v1"
)

func TestNative(t *testing.T) { TestingT(t) }

type NativeSuite struct {
	suite *test.AuthSuite
}

var _ = Suite(&NativeSuite{})

func (s *NativeSuite) SetUpSuite(c *C) {
	utils.InitLoggerForTests()
	PrecalculatedKeysNum = 1
	s.suite = &test.AuthSuite{A: New()}
}

func (s *NativeSuite) TestGenerateKeypairEmptyPass(c *C) {
	s.suite.GenerateKeypairEmptyPass(c)
}

func (s *NativeSuite) TestGenerateKeypairPass(c *C) {
	s.suite.GenerateKeypairPass(c)
}

func (s *NativeSuite) TestGenerateHostCert(c *C) {
	s.suite.GenerateHostCert(c)
}

func (s *NativeSuite) TestGenerateUserCert(c *C) {
	s.suite.GenerateUserCert(c)
}

// TestBuildPrincipals makes sure that the list of principals for a host
// certificate is correctly built.
//   * If the node has role admin, then only the host ID should be listed
//     in the principals field.
//   * If only a host ID is provided, don't include a empty node name
//     this is for backward compatibility.
//   * If both host ID and node name are given, then both should be included
//     on the certificate.
//   * If the host ID and node name are the same, only list one.
func (s *NativeSuite) TestBuildPrincipals(c *C) {
	caPrivateKey, _, err := s.suite.A.GenerateKeyPair("")
	c.Assert(err, IsNil)

	_, hostPublicKey, err := s.suite.A.GenerateKeyPair("")
	c.Assert(err, IsNil)

	tests := []struct {
		inHostID           string
		inNodeName         string
		inClusterName      string
		inRoles            teleport.Roles
		outValidPrincipals []string
	}{
		// 0 - admin role
		{
			"00000000-0000-0000-0000-000000000000",
			"auth",
			"example.com",
			teleport.Roles{teleport.RoleAdmin},
			[]string{"00000000-0000-0000-0000-000000000000"},
		},
		// 1 - backward compatibility
		{
			"11111111-1111-1111-1111-111111111111",
			"",
			"example.com",
			teleport.Roles{teleport.RoleNode},
			[]string{"11111111-1111-1111-1111-111111111111.example.com"},
		},
		// 2 - dual principals
		{
			"22222222-2222-2222-2222-222222222222",
			"proxy",
			"example.com",
			teleport.Roles{teleport.RoleProxy},
			[]string{"22222222-2222-2222-2222-222222222222.example.com", "proxy.example.com"},
		},
		// 3 - deduplicate principals
		{
			"33333333-3333-3333-3333-333333333333",
			"33333333-3333-3333-3333-333333333333",
			"example.com",
			teleport.Roles{teleport.RoleProxy},
			[]string{"33333333-3333-3333-3333-333333333333.example.com"},
		},
	}

	// run tests
	for _, tt := range tests {
		hostCertificateBytes, err := s.suite.A.GenerateHostCert(
			services.CertParams{
				PrivateCASigningKey: caPrivateKey,
				PublicHostKey:       hostPublicKey,
				HostID:              tt.inHostID,
				NodeName:            tt.inNodeName,
				ClusterName:         tt.inClusterName,
				Roles:               tt.inRoles,
				TTL:                 time.Hour,
			})
		c.Assert(err, IsNil)

		publicKey, _, _, _, err := ssh.ParseAuthorizedKey(hostCertificateBytes)
		c.Assert(err, IsNil)

		hostCertificate, ok := publicKey.(*ssh.Certificate)
		c.Assert(ok, Equals, true)

		c.Assert(hostCertificate.ValidPrincipals, HasLen, len(tt.outValidPrincipals))
		c.Assert(hostCertificate.ValidPrincipals, DeepEquals, tt.outValidPrincipals)
	}
}
<commit_msg>fix test<commit_after>/*
Copyright 2017 Gravitational, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package native

import (
	"testing"
	"time"

	"golang.org/x/crypto/ssh"

	"github.com/gravitational/teleport"
	"github.com/gravitational/teleport/lib/auth/test"
	"github.com/gravitational/teleport/lib/services"
	"github.com/gravitational/teleport/lib/utils"

	. "gopkg.in/check.v1"
)

func TestNative(t *testing.T) { TestingT(t) }

type NativeSuite struct {
	suite *test.AuthSuite
}

var _ = Suite(&NativeSuite{})

func (s *NativeSuite) SetUpSuite(c *C) {
	utils.InitLoggerForTests()
	PrecalculatedKeysNum = 1
	s.suite = &test.AuthSuite{A: New()}
}

func (s *NativeSuite) TestGenerateKeypairEmptyPass(c *C) {
	s.suite.GenerateKeypairEmptyPass(c)
}

func (s *NativeSuite) TestGenerateKeypairPass(c *C) {
	s.suite.GenerateKeypairPass(c)
}

func (s *NativeSuite) TestGenerateHostCert(c *C) {
	s.suite.GenerateHostCert(c)
}

func (s *NativeSuite) TestGenerateUserCert(c *C) {
	s.suite.GenerateUserCert(c)
}

// TestBuildPrincipals makes sure that the list of principals for a host
// certificate is correctly built.
//   * If the node has role admin, then only the host ID should be listed
//     in the principals field.
//   * If only a host ID is provided, don't include a empty node name
//     this is for backward compatibility.
//   * If both host ID and node name are given, then both should be included
//     on the certificate.
//   * If the host ID and node name are the same, only list one.
func (s *NativeSuite) TestBuildPrincipals(c *C) {
	caPrivateKey, _, err := s.suite.A.GenerateKeyPair("")
	c.Assert(err, IsNil)

	_, hostPublicKey, err := s.suite.A.GenerateKeyPair("")
	c.Assert(err, IsNil)

	tests := []struct {
		inHostID           string
		inNodeName         string
		inClusterName      string
		inRoles            teleport.Roles
		outValidPrincipals []string
	}{
		// 0 - admin role
		{
			"00000000-0000-0000-0000-000000000000",
			"auth",
			"example.com",
			teleport.Roles{teleport.RoleAdmin},
			[]string{"00000000-0000-0000-0000-000000000000"},
		},
		// 1 - backward compatibility
		{
			"11111111-1111-1111-1111-111111111111",
			"",
			"example.com",
			teleport.Roles{teleport.RoleNode},
			[]string{"11111111-1111-1111-1111-111111111111.example.com"},
		},
		// 2 - dual principals
		{
			"22222222-2222-2222-2222-222222222222",
			"proxy",
			"example.com",
			teleport.Roles{teleport.RoleProxy},
			[]string{"22222222-2222-2222-2222-222222222222.example.com", "proxy.example.com", "proxy"},
		},
		// 3 - deduplicate principals
		{
			"33333333-3333-3333-3333-333333333333",
			"33333333-3333-3333-3333-333333333333",
			"example.com",
			teleport.Roles{teleport.RoleProxy},
			[]string{"33333333-3333-3333-3333-333333333333.example.com", "33333333-3333-3333-3333-333333333333"},
		},
	}

	// run tests
	for _, tt := range tests {
		hostCertificateBytes, err := s.suite.A.GenerateHostCert(
			services.CertParams{
				PrivateCASigningKey: caPrivateKey,
				PublicHostKey:       hostPublicKey,
				HostID:              tt.inHostID,
				NodeName:            tt.inNodeName,
				ClusterName:         tt.inClusterName,
				Roles:               tt.inRoles,
				TTL:                 time.Hour,
			})
		c.Assert(err, IsNil)

		publicKey, _, _, _, err := ssh.ParseAuthorizedKey(hostCertificateBytes)
		c.Assert(err, IsNil)

		hostCertificate, ok := publicKey.(*ssh.Certificate)
		c.Assert(ok, Equals, true)

		c.Assert(hostCertificate.ValidPrincipals, DeepEquals, tt.outValidPrincipals)
	}
}
<|endoftext|>


<commit_before>// Copyright 2012 Apcera Inc. All rights reserved.

package server

import (
	"bufio"
	"encoding/json"
	"fmt"
	"math/rand"
	"net"
	"sync"
	"time"

	"github.com/apcera/gnatsd/hashmap"
	"github.com/apcera/gnatsd/sublist"
)

// The size of the bufio reader/writer on top of the socket.
const defaultBufSize = 32768

type client struct {
	mu   sync.Mutex
	cid  uint64
	opts clientOpts
	conn net.Conn
	bw   *bufio.Writer
	srv  *Server
	subs *hashmap.HashMap
	pcd  map[*client]struct{}
	atmr *time.Timer
	cstats
	parseState
}

func (c *client) String() string {
	return fmt.Sprintf("cid:%d", c.cid)
}

type cstats struct {
	nr int
	nb int
	nm int
}

type subscription struct {
	client  *client
	subject []byte
	queue   []byte
	sid     []byte
	nm      int64
	max     int64
}

type clientOpts struct {
	Verbose       bool   `json:"verbose"`
	Pedantic      bool   `json:"pedantic"`
	SslRequired   bool   `json:"ssl_required"`
	Authorization string `json:"auth_token"`
	Username      string `json:"user"`
	Password      string `json:"pass"`
	Name          string `json:"name"`
}

var defaultOpts = clientOpts{Verbose: true, Pedantic: true}

func init() {
	rand.Seed(time.Now().UnixNano())
}

func (c *client) readLoop() {
	b := make([]byte, defaultBufSize)
	for {
		n, err := c.conn.Read(b)
		if err != nil {
			c.closeConnection()
			return
		}
		if err := c.parse(b[:n]); err != nil {
			Log(err.Error(), clientConnStr(c.conn), c.cid)
			c.closeConnection()
			return
		}
		// Check pending clients for flush.
		for cp, _ := range c.pcd {
			// Flush those in the set
			cp.mu.Lock()
			err := cp.bw.Flush()
			cp.mu.Unlock()
			if err != nil {
				// FIXME, close connection?
				Logf("Error flushing client connection: %v\n", err)
			}
			delete(c.pcd, cp)
		}
	}
}

func (c *client) traceMsg(msg []byte) {
	pm := fmt.Sprintf("Processing msg: %d", c.nm)
	opa := []interface{}{pm, string(c.pa.subject), string(c.pa.reply), string(msg)}
	Trace(logStr(opa), fmt.Sprintf("c: %d", c.cid))
}

func (c *client) traceOp(op string, arg []byte) {
	if !trace {
		return
	}
	opa := []interface{}{fmt.Sprintf("%s OP", op)}
	if arg != nil {
		opa = append(opa, fmt.Sprintf("%s %s", op, string(arg)))
	}
	Trace(logStr(opa), fmt.Sprintf("c: %d", c.cid))
}

func (c *client) processConnect(arg []byte) error {
	c.traceOp("CONNECT", arg)

	// This will be resolved regardless before we exit this func,
	// so we can just clear it here.
	if c.atmr != nil {
		c.atmr.Stop()
		c.atmr = nil
	}

	// FIXME, check err
	if err := json.Unmarshal(arg, &c.opts); err != nil {
		return err
	}
	// Check for Auth
	if c.srv != nil {
		if ok := c.srv.checkAuth(c); !ok {
			c.sendErr("Authorization is Required")
			return fmt.Errorf("Authorization Error")
		}
	}
	if c.opts.Verbose {
		c.sendOK()
	}
	return nil
}

func (c *client) authViolation() {
	c.sendErr("Authorization is Required")
	fmt.Printf("AUTH TIMER EXPIRED!!\n")
	c.closeConnection()
}

func (c *client) sendErr(err string) {
	c.mu.Lock()
	if c.bw != nil {
		c.bw.WriteString(fmt.Sprintf("-ERR '%s'\r\n", err))
		c.pcd[c] = needFlush
	}
	c.mu.Unlock()
}

func (c *client) sendOK() {
	c.mu.Lock()
	c.bw.WriteString("+OK\r\n")
	c.pcd[c] = needFlush
	c.mu.Unlock()
}

func (c *client) processPing() {
	c.traceOp("PING", nil)
	if c.conn == nil {
		return
	}
	c.mu.Lock()
	c.bw.WriteString("PONG\r\n")
	err := c.bw.Flush()
	c.mu.Unlock()
	if err != nil {
		// FIXME, close connection?
		Logf("Error flushing client connection [PING]: %v\n", err)
	}
}

const argsLenMax = 3

func (c *client) processPub(arg []byte) error {
	if trace {
		c.traceOp("PUB", arg)
	}

	// Unroll splitArgs to avoid runtime/heap issues
	a := [argsLenMax][]byte{}
	args := a[:0]
	start := -1
	for i, b := range arg {
		switch b {
		case ' ', '\t', '\r', '\n':
			if start >= 0 {
				args = append(args, arg[start:i])
				start = -1
			}
		default:
			if start < 0 {
				start = i
			}
		}
	}
	if start >= 0 {
		args = append(args, arg[start:])
	}

	switch len(args) {
	case 2:
		c.pa.subject = args[0]
		c.pa.reply = nil
		c.pa.size = parseSize(args[1])
		c.pa.szb = args[1]
	case 3:
		c.pa.subject = args[0]
		c.pa.reply = args[1]
		c.pa.size = parseSize(args[2])
		c.pa.szb = args[2]
	default:
		return fmt.Errorf("processPub Parse Error: '%s'", arg)
	}
	if c.pa.size < 0 {
		return fmt.Errorf("processPub Bad or Missing Size: '%s'", arg)
	}
	if c.opts.Pedantic && !sublist.IsValidLiteralSubject(c.pa.subject) {
		c.sendErr("Invalid Subject")
	}
	return nil
}

func splitArg(arg []byte) [][]byte {
	a := [argsLenMax][]byte{}
	args := a[:0]
	start := -1
	for i, b := range arg {
		switch b {
		case ' ', '\t', '\r', '\n':
			if start >= 0 {
				args = append(args, arg[start:i])
				start = -1
			}
		default:
			if start < 0 {
				start = i
			}
		}
	}
	if start >= 0 {
		args = append(args, arg[start:])
	}
	return args
}

func (c *client) processSub(argo []byte) (err error) {
	c.traceOp("SUB", argo)
	// Copy so we do not reference a potentially large buffer
	arg := make([]byte, len(argo))
	copy(arg, argo)
	args := splitArg(arg)
	sub := &subscription{client: c}
	switch len(args) {
	case 2:
		sub.subject = args[0]
		sub.queue = nil
		sub.sid = args[1]
	case 3:
		sub.subject = args[0]
		sub.queue = args[1]
		sub.sid = args[2]
	default:
		return fmt.Errorf("processSub Parse Error: '%s'", arg)
	}

	c.mu.Lock()
	c.subs.Set(sub.sid, sub)
	if c.srv != nil {
		err = c.srv.sl.Insert(sub.subject, sub)
	}
	c.mu.Unlock()
	if err != nil {
		c.sendErr("Invalid Subject")
	} else if c.opts.Verbose {
		c.sendOK()
	}
	return nil
}

func (c *client) unsubscribe(sub *subscription) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if sub.max > 0 && sub.nm <= sub.max {
		return
	}
	c.traceOp("DELSUB", sub.sid)
	c.subs.Remove(sub.sid)
	if c.srv != nil {
		c.srv.sl.Remove(sub.subject, sub)
	}
}

func (c *client) processUnsub(arg []byte) error {
	c.traceOp("UNSUB", arg)
	args := splitArg(arg)
	var sid []byte
	max := -1

	switch len(args) {
	case 1:
		sid = args[0]
	case 2:
		sid = args[0]
		max = parseSize(args[1])
	default:
		return fmt.Errorf("processUnsub Parse Error: '%s'", arg)
	}
	if sub, ok := (c.subs.Get(sid)).(*subscription); ok {
		if max > 0 {
			sub.max = int64(max)
		}
		c.unsubscribe(sub)
	}
	if c.opts.Verbose {
		c.sendOK()
	}
	return nil
}

func (c *client) msgHeader(mh []byte, sub *subscription) []byte {
	mh = append(mh, sub.sid...)
	mh = append(mh, ' ')
	if c.pa.reply != nil {
		mh = append(mh, c.pa.reply...)
		mh = append(mh, ' ')
	}
	mh = append(mh, c.pa.szb...)
	mh = append(mh, "\r\n"...)
	return mh
}

// Used to treat map as efficient set
type empty struct{}

var needFlush = empty{}

func (c *client) deliverMsg(sub *subscription, mh, msg []byte) {
	if sub.client == nil || sub.client.conn == nil {
		return
	}
	client := sub.client
	client.mu.Lock()
	sub.nm++
	if sub.max > 0 && sub.nm > sub.max {
		client.mu.Unlock()
		client.unsubscribe(sub)
		return
	}
	// Deliver to the client.
	_, err := client.bw.Write(mh)
	if err != nil {
		Logf("Error writing msg header: %v\n", err)
	}
	_, err = client.bw.Write(msg)
	if err != nil {
		Logf("Error writing msg: %v\n", err)
	}
	// FIXME, this is already attached to original message
	_, err = client.bw.WriteString(CR_LF)
	if err != nil {
		Logf("Error writing CRLF: %v\n", err)
	}
	client.mu.Unlock()
	c.pcd[sub.client] = needFlush
}

func (c *client) processMsg(msg []byte) {
	c.nm++
	if trace {
		c.traceMsg(msg)
	}
	if c.srv == nil {
		return
	}
	if c.opts.Verbose {
		c.sendOK()
	}

	scratch := [512]byte{}
	msgh := scratch[:0]

	r := c.srv.sl.Match(c.pa.subject)
	if len(r) <= 0 {
		return
	}

	// msg header
	// FIXME, put MSG into initializer
	msgh = append(msgh, "MSG "...)
	msgh = append(msgh, c.pa.subject...)
	msgh = append(msgh, ' ')
	si := len(msgh)

	var qmap map[string][]*subscription
	var qsubs []*subscription

	for _, v := range r {
		sub := v.(*subscription)
		if sub.queue != nil {
			// FIXME, this can be more efficient
			if qmap == nil {
				qmap = make(map[string][]*subscription)
			}
			//qname := *(*string)(unsafe.Pointer(&sub.queue))
			qname := string(sub.queue)
			qsubs = qmap[qname]
			if qsubs == nil {
				qsubs = make([]*subscription, 0, 4)
			}
			qsubs = append(qsubs, sub)
			qmap[qname] = qsubs
			continue
		}
		mh := c.msgHeader(msgh[:si], sub)
		c.deliverMsg(sub, mh, msg)
	}
	if qmap != nil {
		for _, qsubs := range qmap {
			index := rand.Int() % len(qsubs)
			sub := qsubs[index]
			mh := c.msgHeader(msgh[:si], sub)
			c.deliverMsg(sub, mh, msg)
		}
	}
}

// Lock should be held
func (c *client) clearAuthTimer() {
	if c.atmr == nil {
		return
	}
	c.atmr.Stop()
	c.atmr = nil
}

// Lock should be held
func (c *client) clearConnection() {
	if c.conn == nil {
		return
	}
	c.bw.Flush()
	c.conn.Close()
	c.conn = nil
}

func (c *client) closeConnection() {
	if c.conn == nil {
		return
	}
	Debug("Client connection closed", clientConnStr(c.conn), c.cid)

	c.mu.Lock()
	c.clearAuthTimer()
	c.clearConnection()
	subs := c.subs.All()
	c.mu.Unlock()

	if c.srv != nil {
		for _, s := range subs {
			if sub, ok := s.(*subscription); ok {
				c.srv.sl.Remove(sub.subject, sub)
			}
		}
	}

	/*
		log.Printf("Sublist Stats: %+v\n", c.srv.sl.Stats())
		if c.nr > 0 {
			log.Printf("stats: %d %d %d\n", c.nr, c.nb, c.nm)
			log.Printf("bytes per read: %d\n", c.nb/c.nr)
			log.Printf("msgs per read: %d\n", c.nm/c.nr)
		}
	*/
}
<commit_msg>Added in write deadline logic<commit_after>// Copyright 2012 Apcera Inc. All rights reserved.

package server

import (
	"bufio"
	"encoding/json"
	"fmt"
	"math/rand"
	"net"
	"sync"
	"time"

	"github.com/apcera/gnatsd/hashmap"
	"github.com/apcera/gnatsd/sublist"
)

// The size of the bufio reader/writer on top of the socket.
const defaultBufSize = 32768

type client struct {
	mu   sync.Mutex
	cid  uint64
	opts clientOpts
	conn net.Conn
	bw   *bufio.Writer
	srv  *Server
	subs *hashmap.HashMap
	pcd  map[*client]struct{}
	atmr *time.Timer
	cstats
	parseState
}

func (c *client) String() string {
	return fmt.Sprintf("cid:%d", c.cid)
}

type cstats struct {
	nr int
	nb int
	nm int
}

type subscription struct {
	client  *client
	subject []byte
	queue   []byte
	sid     []byte
	nm      int64
	max     int64
}

type clientOpts struct {
	Verbose       bool   `json:"verbose"`
	Pedantic      bool   `json:"pedantic"`
	SslRequired   bool   `json:"ssl_required"`
	Authorization string `json:"auth_token"`
	Username      string `json:"user"`
	Password      string `json:"pass"`
	Name          string `json:"name"`
}

var defaultOpts = clientOpts{Verbose: true, Pedantic: true}

func init() {
	rand.Seed(time.Now().UnixNano())
}

func (c *client) readLoop() {
	b := make([]byte, defaultBufSize)
	for {
		n, err := c.conn.Read(b)
		if err != nil {
			c.closeConnection()
			return
		}
		if err := c.parse(b[:n]); err != nil {
			Log(err.Error(), clientConnStr(c.conn), c.cid)
			c.closeConnection()
			return
		}
		// Check pending clients for flush.
		for cp, _ := range c.pcd {
			// Flush those in the set
			cp.mu.Lock()
			err := cp.bw.Flush()
			cp.mu.Unlock()
			if err != nil {
				// FIXME, close connection?
				Debugf("Error flushing: %v", err)
			}
			delete(c.pcd, cp)
		}
		// Check to see if we got closed, e.g. slow consumer
		if c.conn == nil {
			return
		}
	}
}

func (c *client) traceMsg(msg []byte) {
	pm := fmt.Sprintf("Processing msg: %d", c.nm)
	opa := []interface{}{pm, string(c.pa.subject), string(c.pa.reply), string(msg)}
	Trace(logStr(opa), fmt.Sprintf("c: %d", c.cid))
}

func (c *client) traceOp(op string, arg []byte) {
	if !trace {
		return
	}
	opa := []interface{}{fmt.Sprintf("%s OP", op)}
	if arg != nil {
		opa = append(opa, fmt.Sprintf("%s %s", op, string(arg)))
	}
	Trace(logStr(opa), fmt.Sprintf("c: %d", c.cid))
}

func (c *client) processConnect(arg []byte) error {
	c.traceOp("CONNECT", arg)

	// This will be resolved regardless before we exit this func,
	// so we can just clear it here.
	if c.atmr != nil {
		c.atmr.Stop()
		c.atmr = nil
	}

	// FIXME, check err
	if err := json.Unmarshal(arg, &c.opts); err != nil {
		return err
	}
	// Check for Auth
	if c.srv != nil {
		if ok := c.srv.checkAuth(c); !ok {
			c.sendErr("Authorization is Required")
			return fmt.Errorf("Authorization Error")
		}
	}
	if c.opts.Verbose {
		c.sendOK()
	}
	return nil
}

func (c *client) authViolation() {
	c.sendErr("Authorization is Required")
	fmt.Printf("AUTH TIMER EXPIRED!!\n")
	c.closeConnection()
}

func (c *client) sendErr(err string) {
	c.mu.Lock()
	if c.bw != nil {
		c.bw.WriteString(fmt.Sprintf("-ERR '%s'\r\n", err))
		c.pcd[c] = needFlush
	}
	c.mu.Unlock()
}

func (c *client) sendOK() {
	c.mu.Lock()
	c.bw.WriteString("+OK\r\n")
	c.pcd[c] = needFlush
	c.mu.Unlock()
}

func (c *client) processPing() {
	c.traceOp("PING", nil)
	if c.conn == nil {
		return
	}
	c.mu.Lock()
	c.bw.WriteString("PONG\r\n")
	err := c.bw.Flush()
	c.mu.Unlock()
	if err != nil {
		// FIXME, close connection?
		Logf("Error flushing client connection [PING]: %v\n", err)
	}
}

const argsLenMax = 3

func (c *client) processPub(arg []byte) error {
	if trace {
		c.traceOp("PUB", arg)
	}

	// Unroll splitArgs to avoid runtime/heap issues
	a := [argsLenMax][]byte{}
	args := a[:0]
	start := -1
	for i, b := range arg {
		switch b {
		case ' ', '\t', '\r', '\n':
			if start >= 0 {
				args = append(args, arg[start:i])
				start = -1
			}
		default:
			if start < 0 {
				start = i
			}
		}
	}
	if start >= 0 {
		args = append(args, arg[start:])
	}

	switch len(args) {
	case 2:
		c.pa.subject = args[0]
		c.pa.reply = nil
		c.pa.size = parseSize(args[1])
		c.pa.szb = args[1]
	case 3:
		c.pa.subject = args[0]
		c.pa.reply = args[1]
		c.pa.size = parseSize(args[2])
		c.pa.szb = args[2]
	default:
		return fmt.Errorf("processPub Parse Error: '%s'", arg)
	}
	if c.pa.size < 0 {
		return fmt.Errorf("processPub Bad or Missing Size: '%s'", arg)
	}
	if c.opts.Pedantic && !sublist.IsValidLiteralSubject(c.pa.subject) {
		c.sendErr("Invalid Subject")
	}
	return nil
}

func splitArg(arg []byte) [][]byte {
	a := [argsLenMax][]byte{}
	args := a[:0]
	start := -1
	for i, b := range arg {
		switch b {
		case ' ', '\t', '\r', '\n':
			if start >= 0 {
				args = append(args, arg[start:i])
				start = -1
			}
		default:
			if start < 0 {
				start = i
			}
		}
	}
	if start >= 0 {
		args = append(args, arg[start:])
	}
	return args
}

func (c *client) processSub(argo []byte) (err error) {
	c.traceOp("SUB", argo)
	// Copy so we do not reference a potentially large buffer
	arg := make([]byte, len(argo))
	copy(arg, argo)
	args := splitArg(arg)
	sub := &subscription{client: c}
	switch len(args) {
	case 2:
		sub.subject = args[0]
		sub.queue = nil
		sub.sid = args[1]
	case 3:
		sub.subject = args[0]
		sub.queue = args[1]
		sub.sid = args[2]
	default:
		return fmt.Errorf("processSub Parse Error: '%s'", arg)
	}

	c.mu.Lock()
	c.subs.Set(sub.sid, sub)
	if c.srv != nil {
		err = c.srv.sl.Insert(sub.subject, sub)
	}
	c.mu.Unlock()
	if err != nil {
		c.sendErr("Invalid Subject")
	} else if c.opts.Verbose {
		c.sendOK()
	}
	return nil
}

func (c *client) unsubscribe(sub *subscription) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if sub.max > 0 && sub.nm <= sub.max {
		return
	}
	c.traceOp("DELSUB", sub.sid)
	c.subs.Remove(sub.sid)
	if c.srv != nil {
		c.srv.sl.Remove(sub.subject, sub)
	}
}

func (c *client) processUnsub(arg []byte) error {
	c.traceOp("UNSUB", arg)
	args := splitArg(arg)
	var sid []byte
	max := -1

	switch len(args) {
	case 1:
		sid = args[0]
	case 2:
		sid = args[0]
		max = parseSize(args[1])
	default:
		return fmt.Errorf("processUnsub Parse Error: '%s'", arg)
	}
	if sub, ok := (c.subs.Get(sid)).(*subscription); ok {
		if max > 0 {
			sub.max = int64(max)
		}
		c.unsubscribe(sub)
	}
	if c.opts.Verbose {
		c.sendOK()
	}
	return nil
}

func (c *client) msgHeader(mh []byte, sub *subscription) []byte {
	mh = append(mh, sub.sid...)
	mh = append(mh, ' ')
	if c.pa.reply != nil {
		mh = append(mh, c.pa.reply...)
		mh = append(mh, ' ')
	}
	mh = append(mh, c.pa.szb...)
	mh = append(mh, "\r\n"...)
	return mh
}

// Used to treat map as efficient set
type empty struct{}

var needFlush = empty{}

func (c *client) deliverMsg(sub *subscription, mh, msg []byte) {
	if sub.client == nil || sub.client.conn == nil {
		return
	}
	client := sub.client
	client.mu.Lock()
	sub.nm++
	if sub.max > 0 && sub.nm > sub.max {
		client.mu.Unlock()
		client.unsubscribe(sub)
		return
	}

	// Check to see if our writes will cause a flush
	// in the underlying bufio. If so limit time we
	// will wait for flush to complete.

	deadlineSet := false
	if client.bw.Available() < (len(mh) + len(msg) + len(CR_LF)) {
		client.conn.SetWriteDeadline(time.Now().Add(DEFAULT_FLUSH_DEADLINE))
		deadlineSet = true
	}

	// Deliver to the client.
	_, err := client.bw.Write(mh)
	if err != nil {
		goto writeErr
	}

	_, err = client.bw.Write(msg)
	if err != nil {
		goto writeErr
	}

	// FIXME, this is already attached to original message
	_, err = client.bw.WriteString(CR_LF)
	if err != nil {
		goto writeErr
	}

	if deadlineSet {
		client.conn.SetWriteDeadline(time.Time{})
	}

	client.mu.Unlock()
	c.pcd[client] = needFlush
	return

writeErr:
	if deadlineSet {
		client.conn.SetWriteDeadline(time.Time{})
	}
	client.mu.Unlock()

	if ne, ok := err.(net.Error); ok && ne.Timeout() {
		// FIXME: SlowConsumer logic
		Log("Slow Consumer Detected", clientConnStr(client.conn), client.cid)
		client.closeConnection()
	} else {
		Debugf("Error writing msg: %v", err)
	}
}

func (c *client) processMsg(msg []byte) {
	c.nm++
	if trace {
		c.traceMsg(msg)
	}
	if c.srv == nil {
		return
	}
	if c.opts.Verbose {
		c.sendOK()
	}

	scratch := [512]byte{}
	msgh := scratch[:0]

	r := c.srv.sl.Match(c.pa.subject)
	if len(r) <= 0 {
		return
	}

	// msg header
	// FIXME, put MSG into initializer
	msgh = append(msgh, "MSG "...)
	msgh = append(msgh, c.pa.subject...)
	msgh = append(msgh, ' ')
	si := len(msgh)

	var qmap map[string][]*subscription
	var qsubs []*subscription

	for _, v := range r {
		sub := v.(*subscription)
		if sub.queue != nil {
			// FIXME, this can be more efficient
			if qmap == nil {
				qmap = make(map[string][]*subscription)
			}
			//qname := *(*string)(unsafe.Pointer(&sub.queue))
			qname := string(sub.queue)
			qsubs = qmap[qname]
			if qsubs == nil {
				qsubs = make([]*subscription, 0, 4)
			}
			qsubs = append(qsubs, sub)
			qmap[qname] = qsubs
			continue
		}
		mh := c.msgHeader(msgh[:si], sub)
		c.deliverMsg(sub, mh, msg)
	}
	if qmap != nil {
		for _, qsubs := range qmap {
			index := rand.Int() % len(qsubs)
			sub := qsubs[index]
			mh := c.msgHeader(msgh[:si], sub)
			c.deliverMsg(sub, mh, msg)
		}
	}
}

// Lock should be held
func (c *client) clearAuthTimer() {
	if c.atmr == nil {
		return
	}
	c.atmr.Stop()
	c.atmr = nil
}

// Lock should be held
func (c *client) clearConnection() {
	if c.conn == nil {
		return
	}
	c.bw.Flush()
	c.conn.Close()
	c.conn = nil
}

func (c *client) closeConnection() {
	if c.conn == nil {
		return
	}
	Debug("Client connection closed", clientConnStr(c.conn), c.cid)

	c.mu.Lock()
	c.clearAuthTimer()
	c.clearConnection()
	subs := c.subs.All()
	c.mu.Unlock()

	if c.srv != nil {
		for _, s := range subs {
			if sub, ok := s.(*subscription); ok {
				c.srv.sl.Remove(sub.subject, sub)
			}
		}
	}

	/*
		log.Printf("Sublist Stats: %+v\n", c.srv.sl.Stats())
		if c.nr > 0 {
			log.Printf("stats: %d %d %d\n", c.nr, c.nb, c.nm)
			log.Printf("bytes per read: %d\n", c.nb/c.nr)
			log.Printf("msgs per read: %d\n", c.nm/c.nr)
		}
	*/
}
<|endoftext|>