<commit_before>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package sync

import (
	"errors"
	"fmt"
	"testing"
	"time"

	"github.com/golang/mock/gomock"
	"github.com/stretchr/testify/assert"
	"go.uber.org/yarpc/internal/testtime"
)

func TestLifecycleOnce(t *testing.T) {
	type testStruct struct {
		msg string

		// A list of actions that will be applied on the LifecycleOnce
		actions []LifecycleAction

		// expected state at the end of the actions
		expectedFinalState LifecycleState
	}
	tests := []testStruct{
		{
			msg:                "setup",
			expectedFinalState: Idle,
		},
		{
			msg: "Start",
			actions: []LifecycleAction{
				StartAction{ExpectedState: Running},
			},
			expectedFinalState: Running,
		},
		{
			msg: "Start and Started",
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StartAction{ExpectedState: Running},
						Actions{
							WaitForStartAction,
							GetStateAction{ExpectedState: Running},
						},
					},
				},
			},
			expectedFinalState: Running,
		},
		{
			msg: "Stop",
			actions: []LifecycleAction{
				StartAction{ExpectedState: Running},
				StopAction{ExpectedState: Stopped},
			},
			expectedFinalState: Stopped,
		},
		{
			msg: "Stop and Stopped",
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StopAction{ExpectedState: Stopped},
						WaitForStoppingAction,
						WaitForStopAction,
					},
				},
			},
			expectedFinalState: Stopped,
		},
		{
			msg: "Error and Stopped",
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StartAction{
							Err:           fmt.Errorf("abort"),
							ExpectedErr:   fmt.Errorf("abort"),
							ExpectedState: Errored,
						},
						Actions{
							WaitForStoppingAction,
							GetStateAction{ExpectedState: Errored},
						},
						Actions{
							WaitForStopAction,
							GetStateAction{ExpectedState: Errored},
						},
					},
				},
			},
			expectedFinalState: Errored,
		},
		{
			msg: "Start, Stop, and Stopped",
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						Actions{
							StartAction{ExpectedState: Running},
							StopAction{ExpectedState: Stopped},
						},
						WaitForStoppingAction,
						WaitForStopAction,
					},
				},
			},
			expectedFinalState: Stopped,
		},
		{
			msg: "Starting",
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StartAction{ExpectedState: Running, Wait: 20 * testtime.Millisecond},
						GetStateAction{ExpectedState: Starting},
					},
					Wait: 10 * testtime.Millisecond,
				},
			},
			expectedFinalState: Running,
		},
		{
			msg: "Stopping",
			actions: []LifecycleAction{
				StartAction{ExpectedState: Running},
				ConcurrentAction{
					Actions: []LifecycleAction{
						StopAction{ExpectedState: Stopped, Wait: 20 * testtime.Millisecond},
						GetStateAction{ExpectedState: Stopping},
						GetStateAction{ExpectedState: Stopped},
					},
					Wait: 10 * testtime.Millisecond,
				},
			},
			expectedFinalState: Stopped,
		},
		{
			msg: "Delayed stop, wait for stopping",
			// The purpose of this test is to verify that the stopping state is
			// occupied for the duration of the Stop() call.
			//
			// Timeline:
			// - 0ms Idle
			// - 0ms Starting
			// - 0–20ms Running
			// - 20–40ms Stopping
			// - 40–ms Stopped
			//
			// First group:
			// - 0ms Start (0ms)
			// - 0ms Wait (20ms)
			// - 20ms Stop (20ms)
			//
			// Second group:
			// - 0ms Wait until Stopping
			// - 20ms Resume
			// - 20ms Wait (10ms)
			// - 30ms Expect Stopping state
			// - 30ms Wait (20ms)
			// - 50ms Expect Stopped state
			actions: []LifecycleAction{
				StartAction{ExpectedState: Running},
				ConcurrentAction{
					Actions: []LifecycleAction{
						Actions{
							StartAction{ExpectedState: Running},
							WaitAction(20 * testtime.Millisecond),
							StopAction{ExpectedState: Stopped, Wait: 20 * testtime.Millisecond},
						},
						Actions{
							WaitForStoppingAction,
							WaitAction(10 * testtime.Millisecond),
							ExactStateAction{ExpectedState: Stopping},
							WaitAction(20 * testtime.Millisecond),
							GetStateAction{ExpectedState: Stopped},
						},
					},
				},
			},
			expectedFinalState: Stopped,
		},
		{
			msg: "Start assure only called once and propagates the same error",
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StartAction{
							Wait:          40 * testtime.Millisecond,
							Err:           errors.New("expected error"),
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
						StartAction{
							Err:           errors.New("not an expected error 1"),
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
						StartAction{
							Wait:          40 * testtime.Millisecond,
							Err:           errors.New("not an expected error 2"),
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
					},
					Wait: 10 * testtime.Millisecond,
				},
			},
			expectedFinalState: Errored,
		},
		{
			msg: "Successful Start followed by failed Stop",
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StartAction{
							Wait:          10 * testtime.Millisecond,
							ExpectedState: Running,
						},
						StopAction{
							Wait:          10 * testtime.Millisecond,
							Err:           errors.New("expected error"),
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
						StartAction{
							Err:           errors.New("not expected error 2"),
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
						StopAction{
							Err:           errors.New("not expected error 2"),
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
					},
					Wait: 30 * testtime.Millisecond,
				},
			},
			expectedFinalState: Errored,
		},
		{
			msg: "Stop assure only called once and returns the same error",
			actions: []LifecycleAction{
				StartAction{
					ExpectedState: Running,
				},
				ConcurrentAction{
					Actions: []LifecycleAction{
						StopAction{
							Wait:          40 * testtime.Millisecond,
							Err:           errors.New("expected error"),
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
						StopAction{
							Wait:          40 * testtime.Millisecond,
							Err:           errors.New("not an expected error 1"),
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
						StopAction{
							Wait:          40 * testtime.Millisecond,
							Err:           errors.New("not an expected error 2"),
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
					},
					Wait: 10 * testtime.Millisecond,
				},
			},
			expectedFinalState: Errored,
		},
		{
			msg: "Stop before start goes directly to 'stopped'",
			actions: []LifecycleAction{
				StopAction{
					ExpectedState: Stopped,
				},
			},
			expectedFinalState: Stopped,
		},
		{
			msg: "Pre-empting start after stop",
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StopAction{
							Wait:          10 * testtime.Millisecond,
							ExpectedState: Stopped,
						},
						StartAction{
							Err:           fmt.Errorf("start action should not run"),
							Wait:          500 * time.Second,
							ExpectedState: Stopped,
						},
					},
					Wait: 20 * testtime.Millisecond,
				},
			},
			expectedFinalState: Stopped,
		},
		{
			msg: "Overlapping stop after start",
			// ms: timeline
			// 00: 0: start..............starting
			// 10: |  1. stop
			// 50: X..|..................running
			//        |..................stopping
			// 60:    X..................stopped
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StartAction{
							Wait:          50 * testtime.Millisecond,
							ExpectedState: Running,
						},
						StopAction{
							Wait:          10 * testtime.Millisecond,
							ExpectedState: Stopped,
						},
					},
					Wait: 10 * testtime.Millisecond,
				},
			},
			expectedFinalState: Stopped,
		},
		{
			msg: "Overlapping stop after start error",
			// ms: timeline
			// 00: 0: start..............starting
			// 10: |  1. stop............stopping
			// 50: X  X..................errored
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StartAction{
							Wait:          50 * testtime.Millisecond,
							Err:           errors.New("expected error"),
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
						StopAction{
							Wait:          10 * testtime.Millisecond,
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
					},
					Wait: 10 * testtime.Millisecond,
				},
			},
			expectedFinalState: Errored,
		},
		{
			msg: "Overlapping start after stop",
			// ms: timeline
			// 00:  0: start.............starting
			// 10:  |
			// 20:  |  1: stop
			// 30:  X  -.................running
			// 30+Δ:   |.................stopping
			// 40:     |  2: start
			// 40+Δ:   |  X
			// 50:     |
			// 60:     |
			// 70:     X.................stopped
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StartAction{
							Wait:          30 * testtime.Millisecond,
							ExpectedState: Running,
						},
						StopAction{
							Wait:          30 * testtime.Millisecond,
							ExpectedState: Stopped,
						},
						StartAction{
							Err:           fmt.Errorf("start action should not run"),
							ExpectedState: Stopping,
						},
					},
					Wait: 20 * testtime.Millisecond,
				},
			},
			expectedFinalState: Stopped,
		},
		{
			msg: "Start completes before overlapping stop completes",
			// ms: timeline
			// 00:   0: start............starting
			// 10:   |  1: start
			// 20:   |  -  2: stop
			// 30:   X  -  -  3: start...running
			// 30+Δ:    X  |  X..........stopping
			// 40:         |     4: stop
			//             |     -
			//             |     -
			// 60:         X     X.......stopped
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StartAction{
							Wait:          30 * testtime.Millisecond,
							ExpectedState: Running,
						},
						StartAction{
							Err:           fmt.Errorf("start action should not run"),
							ExpectedState: Running,
						},
						StopAction{
							Wait:          40 * testtime.Millisecond,
							ExpectedState: Stopped,
						},
						StartAction{
							Err:           fmt.Errorf("start action should not run"),
							ExpectedState: Running,
						},
						StopAction{
							Err:           fmt.Errorf("stop action should not run"),
							ExpectedState: Stopped,
						},
					},
					Wait: 10 * testtime.Millisecond,
				},
			},
			expectedFinalState: Stopped,
		},
	}

	for _, tt := range tests {
		t.Run(tt.msg, func(t *testing.T) {
			mockCtrl := gomock.NewController(t)
			defer mockCtrl.Finish()

			once := Once()
			ApplyLifecycleActions(t, once, tt.actions)

			assert.Equal(t, tt.expectedFinalState, once.LifecycleState())
		})
	}
}

// TestStopping verifies that a lifecycle object can spawn a goroutine and wait
// for that goroutine to exit in its stopping state.  The goroutine must wrap
// up its work when it detects that the lifecycle has begun stopping.  If it
// waited for the stopped channel, the stop callback would deadlock.
func TestStopping(t *testing.T) {
	l := Once()
	l.Start(nil)

	done := make(chan struct{})
	go func() {
		select {
		case <-l.Stopping():
		case <-time.After(time.Second):
			assert.Fail(t, "deadlock")
		}
		close(done)
	}()

	l.Stop(func() error {
		select {
		case <-done:
		case <-time.After(time.Second):
			assert.Fail(t, "deadlock")
		}
		return nil
	})
}
<commit_msg>make lifecycle test less flaky (#1153)<commit_after>// Copyright (c) 2017 Uber Technologies, Inc.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

package sync

import (
	"errors"
	"fmt"
	"testing"
	"time"

	"github.com/golang/mock/gomock"
	"github.com/stretchr/testify/assert"
	"go.uber.org/yarpc/internal/testtime"
)

func TestLifecycleOnce(t *testing.T) {
	type testStruct struct {
		msg string

		// A list of actions that will be applied on the LifecycleOnce
		actions []LifecycleAction

		// expected state at the end of the actions
		expectedFinalState LifecycleState
	}
	tests := []testStruct{
		{
			msg:                "setup",
			expectedFinalState: Idle,
		},
		{
			msg: "Start",
			actions: []LifecycleAction{
				StartAction{ExpectedState: Running},
			},
			expectedFinalState: Running,
		},
		{
			msg: "Start and Started",
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StartAction{ExpectedState: Running},
						Actions{
							WaitForStartAction,
							GetStateAction{ExpectedState: Running},
						},
					},
				},
			},
			expectedFinalState: Running,
		},
		{
			msg: "Stop",
			actions: []LifecycleAction{
				StartAction{ExpectedState: Running},
				StopAction{ExpectedState: Stopped},
			},
			expectedFinalState: Stopped,
		},
		{
			msg: "Stop and Stopped",
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StopAction{ExpectedState: Stopped},
						WaitForStoppingAction,
						WaitForStopAction,
					},
				},
			},
			expectedFinalState: Stopped,
		},
		{
			msg: "Error and Stopped",
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StartAction{
							Err:           fmt.Errorf("abort"),
							ExpectedErr:   fmt.Errorf("abort"),
							ExpectedState: Errored,
						},
						Actions{
							WaitForStoppingAction,
							GetStateAction{ExpectedState: Errored},
						},
						Actions{
							WaitForStopAction,
							GetStateAction{ExpectedState: Errored},
						},
					},
				},
			},
			expectedFinalState: Errored,
		},
		{
			msg: "Start, Stop, and Stopped",
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						Actions{
							StartAction{ExpectedState: Running},
							StopAction{ExpectedState: Stopped},
						},
						WaitForStoppingAction,
						WaitForStopAction,
					},
				},
			},
			expectedFinalState: Stopped,
		},
		{
			msg: "Starting",
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StartAction{ExpectedState: Running, Wait: 20 * testtime.Millisecond},
						GetStateAction{ExpectedState: Starting},
					},
					Wait: 10 * testtime.Millisecond,
				},
			},
			expectedFinalState: Running,
		},
		{
			msg: "Stopping",
			actions: []LifecycleAction{
				StartAction{ExpectedState: Running},
				ConcurrentAction{
					Actions: []LifecycleAction{
						StopAction{ExpectedState: Stopped, Wait: 50 * testtime.Millisecond},
						GetStateAction{ExpectedState: Stopping},
					},
					Wait: 10 * testtime.Millisecond,
				},
			},
			expectedFinalState: Stopped,
		},
		{
			msg: "Delayed stop, wait for stopping",
			// The purpose of this test is to verify that the stopping state is
			// occupied for the duration of the Stop() call.
			//
			// Timeline:
			// - 0ms Idle
			// - 0ms Starting
			// - 0–20ms Running
			// - 20–40ms Stopping
			// - 40–ms Stopped
			//
			// First group:
			// - 0ms Start (0ms)
			// - 0ms Wait (20ms)
			// - 20ms Stop (20ms)
			//
			// Second group:
			// - 0ms Wait until Stopping
			// - 20ms Resume
			// - 20ms Wait (10ms)
			// - 30ms Expect Stopping state
			// - 30ms Wait (20ms)
			// - 50ms Expect Stopped state
			actions: []LifecycleAction{
				StartAction{ExpectedState: Running},
				ConcurrentAction{
					Actions: []LifecycleAction{
						Actions{
							StartAction{ExpectedState: Running},
							WaitAction(20 * testtime.Millisecond),
							StopAction{ExpectedState: Stopped, Wait: 20 * testtime.Millisecond},
						},
						Actions{
							WaitForStoppingAction,
							WaitAction(10 * testtime.Millisecond),
							ExactStateAction{ExpectedState: Stopping},
							WaitAction(20 * testtime.Millisecond),
							GetStateAction{ExpectedState: Stopped},
						},
					},
				},
			},
			expectedFinalState: Stopped,
		},
		{
			msg: "Start assure only called once and propagates the same error",
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StartAction{
							Wait:          40 * testtime.Millisecond,
							Err:           errors.New("expected error"),
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
						StartAction{
							Err:           errors.New("not an expected error 1"),
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
						StartAction{
							Wait:          40 * testtime.Millisecond,
							Err:           errors.New("not an expected error 2"),
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
					},
					Wait: 10 * testtime.Millisecond,
				},
			},
			expectedFinalState: Errored,
		},
		{
			msg: "Successful Start followed by failed Stop",
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StartAction{
							Wait:          10 * testtime.Millisecond,
							ExpectedState: Running,
						},
						StopAction{
							Wait:          10 * testtime.Millisecond,
							Err:           errors.New("expected error"),
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
						StartAction{
							Err:           errors.New("not expected error 2"),
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
						StopAction{
							Err:           errors.New("not expected error 2"),
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
					},
					Wait: 30 * testtime.Millisecond,
				},
			},
			expectedFinalState: Errored,
		},
		{
			msg: "Stop assure only called once and returns the same error",
			actions: []LifecycleAction{
				StartAction{
					ExpectedState: Running,
				},
				ConcurrentAction{
					Actions: []LifecycleAction{
						StopAction{
							Wait:          40 * testtime.Millisecond,
							Err:           errors.New("expected error"),
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
						StopAction{
							Wait:          40 * testtime.Millisecond,
							Err:           errors.New("not an expected error 1"),
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
						StopAction{
							Wait:          40 * testtime.Millisecond,
							Err:           errors.New("not an expected error 2"),
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
					},
					Wait: 10 * testtime.Millisecond,
				},
			},
			expectedFinalState: Errored,
		},
		{
			msg: "Stop before start goes directly to 'stopped'",
			actions: []LifecycleAction{
				StopAction{
					ExpectedState: Stopped,
				},
			},
			expectedFinalState: Stopped,
		},
		{
			msg: "Pre-empting start after stop",
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StopAction{
							Wait:          10 * testtime.Millisecond,
							ExpectedState: Stopped,
						},
						StartAction{
							Err:           fmt.Errorf("start action should not run"),
							Wait:          500 * time.Second,
							ExpectedState: Stopped,
						},
					},
					Wait: 20 * testtime.Millisecond,
				},
			},
			expectedFinalState: Stopped,
		},
		{
			msg: "Overlapping stop after start",
			// ms: timeline
			// 00: 0: start..............starting
			// 10: |  1. stop
			// 50: X..|..................running
			//        |..................stopping
			// 60:    X..................stopped
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StartAction{
							Wait:          50 * testtime.Millisecond,
							ExpectedState: Running,
						},
						StopAction{
							Wait:          10 * testtime.Millisecond,
							ExpectedState: Stopped,
						},
					},
					Wait: 10 * testtime.Millisecond,
				},
			},
			expectedFinalState: Stopped,
		},
		{
			msg: "Overlapping stop after start error",
			// ms: timeline
			// 00: 0: start..............starting
			// 10: |  1. stop............stopping
			// 50: X  X..................errored
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StartAction{
							Wait:          50 * testtime.Millisecond,
							Err:           errors.New("expected error"),
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
						StopAction{
							Wait:          10 * testtime.Millisecond,
							ExpectedState: Errored,
							ExpectedErr:   errors.New("expected error"),
						},
					},
					Wait: 10 * testtime.Millisecond,
				},
			},
			expectedFinalState: Errored,
		},
		{
			msg: "Overlapping start after stop",
			// ms: timeline
			// 00:  0: start.............starting
			// 10:  |
			// 20:  |  1: stop
			// 30:  X  -.................running
			// 30+Δ:   |.................stopping
			// 40:     |  2: start
			// 40+Δ:   |  X
			// 50:     |
			// 60:     |
			// 70:     X.................stopped
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StartAction{
							Wait:          30 * testtime.Millisecond,
							ExpectedState: Running,
						},
						StopAction{
							Wait:          30 * testtime.Millisecond,
							ExpectedState: Stopped,
						},
						StartAction{
							Err:           fmt.Errorf("start action should not run"),
							ExpectedState: Stopping,
						},
					},
					Wait: 20 * testtime.Millisecond,
				},
			},
			expectedFinalState: Stopped,
		},
		{
			msg: "Start completes before overlapping stop completes",
			// ms: timeline
			// 00:   0: start............starting
			// 10:   |  1: start
			// 20:   |  -  2: stop
			// 30:   X  -  -  3: start...running
			// 30+Δ:    X  |  X..........stopping
			// 40:         |     4: stop
			//             |     -
			//             |     -
			// 60:         X     X.......stopped
			actions: []LifecycleAction{
				ConcurrentAction{
					Actions: []LifecycleAction{
						StartAction{
							Wait:          30 * testtime.Millisecond,
							ExpectedState: Running,
						},
						StartAction{
							Err:           fmt.Errorf("start action should not run"),
							ExpectedState: Running,
						},
						StopAction{
							Wait:          40 * testtime.Millisecond,
							ExpectedState: Stopped,
						},
						StartAction{
							Err:           fmt.Errorf("start action should not run"),
							ExpectedState: Running,
						},
						StopAction{
							Err:           fmt.Errorf("stop action should not run"),
							ExpectedState: Stopped,
						},
					},
					Wait: 10 * testtime.Millisecond,
				},
			},
			expectedFinalState: Stopped,
		},
	}

	for _, tt := range tests {
		t.Run(tt.msg, func(t *testing.T) {
			mockCtrl := gomock.NewController(t)
			defer mockCtrl.Finish()

			once := Once()
			ApplyLifecycleActions(t, once, tt.actions)

			assert.Equal(t, tt.expectedFinalState, once.LifecycleState())
		})
	}
}

// TestStopping verifies that a lifecycle object can spawn a goroutine and wait
// for that goroutine to exit in its stopping state.  The goroutine must wrap
// up its work when it detects that the lifecycle has begun stopping.  If it
// waited for the stopped channel, the stop callback would deadlock.
func TestStopping(t *testing.T) {
	l := Once()
	l.Start(nil)

	done := make(chan struct{})
	go func() {
		select {
		case <-l.Stopping():
		case <-time.After(time.Second):
			assert.Fail(t, "deadlock")
		}
		close(done)
	}()

	l.Stop(func() error {
		select {
		case <-done:
		case <-time.After(time.Second):
			assert.Fail(t, "deadlock")
		}
		return nil
	})
}
<|endoftext|>


<commit_before>package builtins

import (
	"os"
	"path/filepath"
)

type file struct {
	valueStub
}

func NewFileClass() Value {
	f := &file{}
	f.initialize()
	f.AddMethod(NewMethod("expand_path", func(args ...Value) (Value, error) {
		arg1 := args[0].(*StringValue).String()

		if arg1[0] == '~' {
			arg1 = os.Getenv("HOME") + arg1[1:]
		}

		var path string
		if len(args) == 2 {
			path = filepath.Join(args[1].(*StringValue).String(), arg1)
		} else {
			path, _ = filepath.Abs(arg1)
		}

		return NewString(path), nil
	}))

	return f
}

func (file *file) String() string {
	return "File"
}

func (file *file) New() Value {
	return file
}
<commit_msg>File class should have "Class" as its superclass<commit_after>package builtins

import (
	"os"
	"path/filepath"
)

type fileClass struct {
	valueStub
}

func NewFileClass() Value {
	f := &fileClass{}
	f.initialize()
	f.class = NewClassValue() // FIXME: this should be a global reference
	f.AddMethod(NewMethod("expand_path", func(args ...Value) (Value, error) {
		arg1 := args[0].(*StringValue).String()

		if arg1[0] == '~' {
			arg1 = os.Getenv("HOME") + arg1[1:]
		}

		var path string
		if len(args) == 2 {
			path = filepath.Join(args[1].(*StringValue).String(), arg1)
		} else {
			path, _ = filepath.Abs(arg1)
		}

		return NewString(path), nil
	}))

	return f
}

func (file *fileClass) String() string {
	return "File"
}

func (file *fileClass) New() Value {
	return file
}
<|endoftext|>


<commit_before>package git_comment

import (
	"errors"
	"fmt"
	git "gopkg.in/libgit2/git2go.v22"
)

const (
	authorNotFoundError  = "No name or email found in git config for commenting"
	commitNotFoundError  = "Commit not found"
	commentNotFoundError = "Comment not found"
	userNameKey          = "user.name"
	userEmailKey         = "user.email"
	headCommit           = "HEAD"
)

// Create a new comment on a commit, optionally with a file and line
func CreateComment(repoPath string, commit *string, fileRef *FileRef, message string) (*string, error) {
	var hash = commit
	author, cErr := author(repoPath)
	if cErr != nil {
		return nil, cErr
	}
	if *hash == "HEAD" || hash == nil {
		head, hErr := head(repoPath)
		if hErr != nil {
			return nil, hErr
		}
		hash = head
	} else {
		// check if commit hash exists
	}
	comment, err := NewComment(message, *hash, fileRef, author)
	if err != nil {
		return nil, err
	}
	fmt.Println(comment.Serialize())
	writeErr := writeCommentToDisk(repoPath, comment)
	if writeErr != nil {
		return nil, writeErr
	}
	return &comment.ID, nil
}

func UpdateComment(repoPath string, ID string, message string) error {
	comment, err := CommentByID(repoPath, ID)
	if err != nil {
		return err
	}
	author, cErr := author(repoPath)
	if cErr != nil {
		return cErr
	}
	comment.Amend(message, author)
	return writeCommentToDisk(repoPath, comment)
}

// Remove a comment from a commit
func DeleteComment(repoPath string, ID string) error {
	comment, err := CommentByID(repoPath, ID)
	if err != nil {
		return err
	}
	comment.Deleted = true
	return writeCommentToDisk(repoPath, comment)
}

// Finds a comment by a given ID
func CommentByID(repoPath string, identifier string) (*Comment, error) {
	return &Comment{}, errors.New(commentNotFoundError)
}

// Finds all comments on a given commit
func CommentsOnCommit(repoPath string, commit string) []*Comment {
	return []*Comment{}
}

// Write git object for a given comment and update the
// comment refs
func writeCommentToDisk(repoPath string, comment *Comment) error {
	return nil
}

func repo(repoPath string) (*git.Repository, error) {
	return git.OpenRepository(repoPath)
}

func author(repoPath string) (*Person, error) {
	repo, err := repo(repoPath)
	if err != nil {
		return nil, err
	}
	config, cErr := repo.Config()
	if cErr != nil {
		return nil, cErr
	}
	name, nErr := config.LookupString(userNameKey)
	email, eErr := config.LookupString(userEmailKey)
	if nErr != nil && eErr != nil {
		return nil, errors.New(authorNotFoundError)
	}
	return &Person{name, email}, nil
}

func head(repoPath string) (*string, error) {
	repo, err := repo(repoPath)
	if err != nil {
		return nil, err
	}
	head, hErr := repo.Head()
	if hErr != nil {
		return nil, hErr
	}
	hash := head.Name()
	return &hash, nil
}
<commit_msg>Cleanup<commit_after>package git_comment

import (
	"errors"
	git "gopkg.in/libgit2/git2go.v22"
)

const (
	authorNotFoundError  = "No name or email found in git config for commenting"
	commitNotFoundError  = "Commit not found"
	commentNotFoundError = "Comment not found"
	userNameKey          = "user.name"
	userEmailKey         = "user.email"
	headCommit           = "HEAD"
)

// Create a new comment on a commit, optionally with a file and line
func CreateComment(repoPath string, commit *string, fileRef *FileRef, message string) (*string, error) {
	var hash = commit
	author, cErr := author(repoPath)
	if cErr != nil {
		return nil, cErr
	}
	if *hash == "HEAD" || hash == nil {
		head, hErr := head(repoPath)
		if hErr != nil {
			return nil, hErr
		}
		hash = head
	} else {
		// check if commit hash exists
	}
	comment, err := NewComment(message, *hash, fileRef, author)
	if err != nil {
		return nil, err
	}
	if writeErr := writeCommentToDisk(repoPath, comment); writeErr != nil {
		return nil, writeErr
	}
	return &comment.ID, nil
}

func UpdateComment(repoPath string, ID string, message string) error {
	comment, err := CommentByID(repoPath, ID)
	if err != nil {
		return err
	}
	author, cErr := author(repoPath)
	if cErr != nil {
		return cErr
	}
	comment.Amend(message, author)
	return writeCommentToDisk(repoPath, comment)
}

// Remove a comment from a commit
func DeleteComment(repoPath string, ID string) error {
	comment, err := CommentByID(repoPath, ID)
	if err != nil {
		return err
	}
	comment.Deleted = true
	return writeCommentToDisk(repoPath, comment)
}

// Finds a comment by a given ID
func CommentByID(repoPath string, identifier string) (*Comment, error) {
	return &Comment{}, errors.New(commentNotFoundError)
}

// Finds all comments on a given commit
func CommentsOnCommit(repoPath string, commit string) []*Comment {
	return []*Comment{}
}

// Write git object for a given comment and update the
// comment refs
func writeCommentToDisk(repoPath string, comment *Comment) error {
	return nil
}

func repo(repoPath string) (*git.Repository, error) {
	return git.OpenRepository(repoPath)
}

func author(repoPath string) (*Person, error) {
	repo, err := repo(repoPath)
	if err != nil {
		return nil, err
	}
	config, cErr := repo.Config()
	if cErr != nil {
		return nil, cErr
	}
	name, nErr := config.LookupString(userNameKey)
	email, eErr := config.LookupString(userEmailKey)
	if nErr != nil && eErr != nil {
		return nil, errors.New(authorNotFoundError)
	}
	return &Person{name, email}, nil
}

func head(repoPath string) (*string, error) {
	repo, err := repo(repoPath)
	if err != nil {
		return nil, err
	}
	head, hErr := repo.Head()
	if hErr != nil {
		return nil, hErr
	}
	hash := head.Name()
	return &hash, nil
}
<|endoftext|>


<commit_before>package ioprogress

import (
	"fmt"
	"io"
	"os"
	"strings"
)

// DrawFunc is the callback type for drawing progress.
type DrawFunc func(int64, int64) error

// DrawTextFormatFunc is a callback used by DrawFuncs that draw text in
// order to format the text into some more human friendly format.
type DrawTextFormatFunc func(int64, int64) string

var defaultDrawFunc DrawFunc

func init() {
	defaultDrawFunc = DrawTerminal(os.Stdout)
}

// DrawTerminal returns a DrawFunc that draws a progress bar to an io.Writer
// that is assumed to be a terminal (and therefore respects carriage returns).
func DrawTerminal(w io.Writer) DrawFunc {
	var maxLength int
	return DrawTerminalf(w, func(progress, total int64) string {
		line := fmt.Sprintf("%d/%d", progress, total)

		// Make sure we pad it to the max length we've ever drawn so that
		// we don't have trailing characters.
		if len(line) < maxLength {
			line = fmt.Sprintf(
				"%s%s", line, strings.Repeat(" ", maxLength-len(line)))
		}
		maxLength = len(line)

		return line
	})
}

// DrawTerminalf returns a DrawFunc that draws a progress bar to an io.Writer
// that is formatted with the given formatting function.
func DrawTerminalf(w io.Writer, f DrawTextFormatFunc) DrawFunc {
	return func(progress, total int64) error {
		if progress == -1 || total == -1 {
			_, err := fmt.Fprintf(w, "\n")
			return err
		}

		_, err := fmt.Fprint(w, f(progress, total)+"\r")
		return err
	}
}

var byteUnits = []string{"B", "KB", "MB", "GB", "TB", "PB"}

// DrawTextFormatBytes is a DrawTextFormatFunc that formats the progress
// and total into human-friendly byte formats.
func DrawTextFormatBytes(progress, total int64) string {
	return fmt.Sprintf("%s/%s", byteUnitStr(progress), byteUnitStr(total))
}

func byteUnitStr(n int64) string {
	var unit string
	size := float64(n)
	for i := 1; i < len(byteUnits); i++ {
		mult := float64(i * 1000)
		if size < mult {
			unit = byteUnits[i-1]
			break
		}

		size = size / mult
	}

	return fmt.Sprintf("%.3g %s", size, unit)
}
<commit_msg>Fix byte formatting<commit_after>package ioprogress

import (
	"fmt"
	"io"
	"os"
	"strings"
)

// DrawFunc is the callback type for drawing progress.
type DrawFunc func(int64, int64) error

// DrawTextFormatFunc is a callback used by DrawFuncs that draw text in
// order to format the text into some more human friendly format.
type DrawTextFormatFunc func(int64, int64) string

var defaultDrawFunc DrawFunc

func init() {
	defaultDrawFunc = DrawTerminal(os.Stdout)
}

// DrawTerminal returns a DrawFunc that draws a progress bar to an io.Writer
// that is assumed to be a terminal (and therefore respects carriage returns).
func DrawTerminal(w io.Writer) DrawFunc {
	var maxLength int
	return DrawTerminalf(w, func(progress, total int64) string {
		line := fmt.Sprintf("%d/%d", progress, total)

		// Make sure we pad it to the max length we've ever drawn so that
		// we don't have trailing characters.
		if len(line) < maxLength {
			line = fmt.Sprintf(
				"%s%s", line, strings.Repeat(" ", maxLength-len(line)))
		}
		maxLength = len(line)

		return line
	})
}

// DrawTerminalf returns a DrawFunc that draws a progress bar to an io.Writer
// that is formatted with the given formatting function.
func DrawTerminalf(w io.Writer, f DrawTextFormatFunc) DrawFunc {
	return func(progress, total int64) error {
		if progress == -1 || total == -1 {
			_, err := fmt.Fprintf(w, "\n")
			return err
		}

		_, err := fmt.Fprint(w, f(progress, total)+"\r")
		return err
	}
}

var byteUnits = []string{"B", "KB", "MB", "GB", "TB", "PB"}

// DrawTextFormatBytes is a DrawTextFormatFunc that formats the progress
// and total into human-friendly byte formats.
func DrawTextFormatBytes(progress, total int64) string {
	return fmt.Sprintf("%s/%s", byteUnitStr(progress), byteUnitStr(total))
}

func byteUnitStr(n int64) string {
	var unit string
	size := float64(n)
	for i := 1; i < len(byteUnits); i++ {
		if size < 1000 {
			unit = byteUnits[i-1]
			break
		}

		size = size / 1000
	}

	return fmt.Sprintf("%.3g %s", size, unit)
}
<|endoftext|>


<commit_before>package main

import (
  "bytes"
  "config"
  "fmt"
  "io"
  _ "io/ioutil"
  "log"
  "net"
  "net/http"
  "net/url"
  "os"
  "os/signal"
  "strings"
  "syscall"
  "time"
)

const (
  CLR_N = "\x1b[0m"
  /* you use codes 30+i to specify foreground color, 40+i to specify background color */
  BLACK   = 0
  RED     = 1
  GREEN   = 2
  YELLO   = 3
  BLUE    = 4
  MAGENTA = 5
  CYAN    = 6
  WHITE   = 7
)

type Server struct {
}

//返回ANSI 控制台颜色格式的字符串
//bc 背景颜色
//fc 前景(文字)颜色
func ansi_color(bc int, fc int, s string) string {
  return fmt.Sprintf("\x1b[%d;%dm%s%s", 40+bc, 30+fc, s, CLR_N)
}

/**
http Header Copay
*/
func header_copy(s http.Header, d *http.Header) {
  for hk, _ := range s {
    d.Set(hk, s.Get(hk))
  }
}

func access_log(w http.ResponseWriter, r *http.Request, query_url string, startTime time.Time) {
  remoteAddr := strings.Split(r.RemoteAddr, ":")[0] //客户端地址

  switch {
  case len(r.Header.Get("X-Real-Ip")) > 0:
    remoteAddr = r.Header.Get("X-Real-Ip")
  case len(r.Header.Get("Remote-Addr")) > 0:
    remoteAddr = r.Header.Get("Remote-Addr")
  case len(r.Header.Get("X-Forwarded-For")) > 0:
    remoteAddr = r.Header.Get("X-Forwarded-For")
  }

  if remoteAddr == "[" || len(remoteAddr) == 0 {
    remoteAddr = "127.0.0.1"
  }

  r.ParseForm()
  var postValues []string
  for k, _ := range r.PostForm {
    postValues = append(postValues, fmt.Sprintf("%s=%s", url.QueryEscape(k), url.QueryEscape(r.FormValue(k))))
  }

  if len(postValues) == 0 {
    postValues = append(postValues, "-")
  }

  content_len := w.Header().Get("Content-Length")
  if len(content_len) == 0 {
    content_len = "-"
  }

  logLine := fmt.Sprintf(`[%s] [%s] S:"%s %s %s F:{%s}" D:"%s F:{%s} %s %s",%0.5fs`,
    ansi_color(WHITE, BLACK, remoteAddr),
    time.Now().Format("2006-01-02 15:04:05.999999999 -0700 MST"),
    r.Method,
    ansi_color(GREEN, BLACK, r.RequestURI),
    r.Proto,
    ansi_color(GREEN, BLACK, strings.Join(postValues, "&")),
    ansi_color(YELLO, BLACK, query_url),
    ansi_color(YELLO, BLACK, strings.Join(postValues, "&")),
    w.Header().Get("Status"),
    content_len,
    time.Now().Sub(startTime).Seconds(),
  )
  g_env.AccessLog.Println(logLine)
}

func parse_querys(r *http.Request) (raw_query []string) {
  r.ParseForm()
  check_key_map := make(map[string]bool)
  check_and_append := func(_k, _v string) {
    if _, ok := check_key_map[_k]; ok {
      return
    }
    check_key_map[_k] = true
    raw_query = append(raw_query, fmt.Sprintf("%s=%s", url.QueryEscape(_k), url.QueryEscape(_v)))
  }

  for k, _ := range r.Form {
    check_and_append(k, r.Form.Get(k))
  }

  if len(r.Referer()) > 0 {
    if uri, err := url.Parse(r.Referer()); err == nil {
      for k, _ := range uri.Query() {
        check_and_append(k, uri.Query().Get(k))
      }
    }
  }
  return
}

/**
获取目标地址
*/
func target_path(r *http.Request, cfg *config.Config) (t string) {
  if len(cfg.TargetPath) > 0 {
    t = cfg.TargetPath
  } else {
    t = r.URL.Path
  }
  return
}

/**
获取目标服务服务器
*/
func target_server(cfg *config.Config) (s string) {
  if len(cfg.TargetServer) > 0 {
    s = cfg.TargetServer
  } else {
    s = g_config.Default.TargetServer
  }
  return
}

func parse_query_key_value(v string) (kv map[string]string) {
  kv = make(map[string]string)
  var _k, _v string
  si := strings.IndexByte(v, '=')
  if si < 0 {
    _k = v
  } else {
    _k = v[0:si]
    if si+1 <= len(v) {
      _v = v[si+1:]
    }
  }
  kv[url.QueryEscape(_k)] = url.QueryEscape(_v)
  return
}

/**
标准的 http 协议，会将 & 和 ;看成参数对的分割符
*/
func parse_query_values(r *http.Request) (vs []map[string]string) {
  for _, v := range strings.Split(r.URL.RawQuery, "&") {
    vs = append(vs, parse_query_key_value(v))
  }
  return
}

/**
获取查询参数并做替换
对GET参数不做标准分割 &; 这两个字符
*/
func __swap_raw_query(r *http.Request, cfg *config.Config) (q string) {
  var tmp_slice []string
  raw_querys := parse_query_values(r)
  append_slict := func(key string, value string) {
    if len(key) == 0 {
      return
    }
    tmp_slice = append(tmp_slice, fmt.Sprintf("%s=%s", key, value))
  }
  if len(cfg.TargetParamNameSwap) == 0 {
    for _, vs := range raw_querys {
      for k1, v1 := range vs {
        append_slict(k1, v1)
      }
    }
    q = strings.Join(tmp_slice, "&")
    return
  }

  for _, vs := range raw_querys {
    for k1, v1 := range vs {
      if k2, ok := cfg.TargetParamNameSwap[k1]; ok {
        append_slict(k2, v1)
      } else {
        append_slict(k1, v1)
      }
    }
  }

  q = strings.Join(tmp_slice, "&")
  return
}

/*
标准http解析库实现
*/
func swap_raw_query(r *http.Request, cfg *config.Config) (q string) {
  var tmp_slice []string
  append_slict := func(key string, value string) {
    tmp_slice = append(tmp_slice, fmt.Sprintf("%s=%s", url.QueryEscape(key), url.QueryEscape(r.URL.Query().Get(value))))
  }
  //如果配置段没有需要交换的参数,则直接返回查询字符串
  if len(cfg.TargetParamNameSwap) == 0 {
    for k, _ := range r.URL.Query() {
      append_slict(k, k)
    }
    q = strings.Join(tmp_slice, "&")
    return
  }

  for k, _ := range r.URL.Query() {
    if v, ok := cfg.TargetParamNameSwap[k]; ok {
      append_slict(v, k)
    } else {
      append_slict(k, k)
    }
  }
  q = strings.Join(tmp_slice, "&")
  return
}

func timeout_dialer(conn_timeout int, rw_timeout int) func(net, addr string) (c net.Conn, err error) {
  return func(netw, addr string) (net.Conn, error) {
    conn, err := net.DialTimeout(netw, addr, time.Duration(conn_timeout)*time.Second)
    if err != nil {
      log.Printf("Failed to connect to [%s]. Timed out after %d seconds\n", addr, rw_timeout)
      return nil, err
    }
    conn.SetDeadline(time.Now().Add(time.Duration(rw_timeout) * time.Second))
    return conn, nil
  }
}

func (s Server) handler_func(w http.ResponseWriter, r *http.Request) {
  var (
    cfg                                  *config.Config
    cfg_err                              *config.ConfigErr
    conntction_timeout, response_timeout int
    req                                  *http.Request
    err                                  error
    raw_query                            []string
  )
  defer func() {
    if re := recover(); re != nil {
      g_env.ErrorLog.Println("Recovered in backendServer:", re)
    }
  }()

  defer r.Body.Close()

  start_at := time.Now()
  raw_query = parse_querys(r)

  if err != nil {
    g_env.ErrorLog.Println(req, err)
    http.Error(w, "Read Body Error.", http.StatusInternalServerError)
    return
  }

  //获取配置文件
  if cfg, cfg_err = g_config.FindBySourcePathAndParams(raw_query, r.URL.Path); cfg_err != nil {
    cfg = g_config.FindByParamsOrSourcePath(raw_query, r.URL.Path)
  }

  if conntction_timeout = cfg.ConnectionTimeout; conntction_timeout <= 0 {
    conntction_timeout = 15
  }

  if response_timeout = cfg.ResponseTimeout; response_timeout <= 0 {
    response_timeout = 120
  }

  transport := http.Transport{
    Dial: timeout_dialer(conntction_timeout, response_timeout),
    ResponseHeaderTimeout: time.Duration(response_timeout) * time.Second,
    DisableCompression:    false,
    DisableKeepAlives:     true,
    MaxIdleConnsPerHost:   2,
  }
  defer transport.CloseIdleConnections()

  client := &http.Client{
    Transport: &transport,
  }

  query_url, _ := url.Parse(target_server(cfg) + target_path(r, cfg) + "?" + swap_raw_query(r, cfg))

  switch r.Method {
  case "GET", "HEAD":
    req, err = http.NewRequest(r.Method, query_url.String(), nil)
  case "POST":
    req, err = http.NewRequest(r.Method, query_url.String(), bytes.NewBufferString(strings.Join(raw_query, "&")))
    req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
  default:
    http.Error(w, "MethodNotAllowed", http.StatusMethodNotAllowed)
    return
  }
  req.Close = true

  header_copy(r.Header, &req.Header)

  if err != nil {
    g_env.ErrorLog.Println(err)
    http.Error(w, err.Error(), http.StatusInternalServerError)
    return
  }

  resp, err := client.Do(req)
  defer resp.Body.Close()

  if err != nil {
    g_env.ErrorLog.Println(req, err)
    http.Error(w, err.Error(), http.StatusInternalServerError)
    return
  }

  for hk, _ := range resp.Header {
    w.Header().Set(hk, resp.Header.Get(hk))
  }
  w.Header().Set("X-Transit-Ver", "0.0.1")
  w.Header().Set("Server", "X-Transit")

  w.WriteHeader(resp.StatusCode)
  io.Copy(w, resp.Body)
  access_log(w, r, query_url.String(), start_at)
}

//TODO
// add unix socket listener

func Run() {
  g_env.ErrorLog.Printf("start@ %s:%d %v \n", g_config.Listen.Host, g_config.Listen.Port, time.Now())
  fmt.Printf("start@ %s:%d %v \n", g_config.Listen.Host, g_config.Listen.Port, time.Now())

  sigchan := make(chan os.Signal, 1)
  signal.Notify(sigchan, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP, syscall.SIGINT, syscall.SIGQUIT)

  server := Server{}
  go func() {
    s := &http.Server{
      Addr:           fmt.Sprintf("%s:%d", g_config.Listen.Host, g_config.Listen.Port),
      Handler:        http.HandlerFunc(server.handler_func),
      ReadTimeout:    120 * time.Second,
      WriteTimeout:   120 * time.Second,
      MaxHeaderBytes: 1 << 20,
    }
    log.Fatal(s.ListenAndServe())
  }()

  <-sigchan
}
<commit_msg>add access begin logger<commit_after>package main

import (
  "bytes"
  "config"
  "crypto/md5"
  "fmt"
  "io"
  _ "io/ioutil"
  "log"
  "math/rand"
  "net"
  "net/http"
  "net/url"
  "os"
  "os/signal"
  "strconv"
  "strings"
  "syscall"
  "time"
)

const (
  CLR_N = "\x1b[0m"
  /* you use codes 30+i to specify foreground color, 40+i to specify background color */
  BLACK   = 0
  RED     = 1
  GREEN   = 2
  YELLO   = 3
  BLUE    = 4
  MAGENTA = 5
  CYAN    = 6
  WHITE   = 7
)

type Server struct {
}

//返回ANSI 控制台颜色格式的字符串
//bc 背景颜色
//fc 前景(文字)颜色
func ansi_color(bc int, fc int, s string) string {
  return fmt.Sprintf("\x1b[%d;%dm%s%s", 40+bc, 30+fc, s, CLR_N)
}

/**
http Header Copay
*/
func header_copy(s http.Header, d *http.Header) {
  for hk, _ := range s {
    d.Set(hk, s.Get(hk))
  }
}

func access_id(as string) string {
  md5 := func(s string) string {
    hash_md5 := md5.New()
    io.WriteString(hash_md5, s)
    return fmt.Sprintf("%x", hash_md5.Sum(nil))
  }
  nano := time.Now().UnixNano()
  rand.Seed(nano)
  rnd_num := rand.Int63()
  fs := md5(md5(as) + md5(strconv.FormatInt(nano, 10)) + md5(strconv.FormatInt(rnd_num, 10)))
  return strings.ToUpper(fmt.Sprintf("%s.%s", fs[0:4], fs[28:]))
}

func access_ip(r *http.Request) string {
  remote_addr := strings.Split(r.RemoteAddr, ":")[0] //客户端地址
  switch {
  case len(r.Header.Get("X-Real-Ip")) > 0:
    remote_addr = r.Header.Get("X-Real-Ip")
  case len(r.Header.Get("Remote-Addr")) > 0:
    remote_addr = r.Header.Get("Remote-Addr")
  case len(r.Header.Get("X-Forwarded-For")) > 0:
    remote_addr = r.Header.Get("X-Forwarded-For")
  }

  if remote_addr == "[" || len(remote_addr) == 0 {
    remote_addr = "127.0.0.1"
  }
  return remote_addr
}

func post_values(r *http.Request) (ps []string) {
  r.ParseForm()
  for k, _ := range r.PostForm {
    ps = append(ps, fmt.Sprintf("%s=%s", url.QueryEscape(k), url.QueryEscape(r.FormValue(k))))
  }
  return
}

//access begin logger
func access_log_begin(aid string, r *http.Request, query_url string) {
  post_values := post_values(r)

  if len(post_values) == 0 {
    post_values = append(post_values, "-")
  }
  log_line := fmt.Sprintf(`Begin [%s] "%s" [%s] S:"%s %s %s F:{%s}" D:"%s F:{%s}"`,
    ansi_color(WHITE, BLACK, access_ip(r)),
    ansi_color(WHITE, RED, aid),
    time.Now().Format("2006-01-02 15:04:05.999999999 -0700 MST"),
    r.Method,
    ansi_color(GREEN, BLACK, r.RequestURI),
    r.Proto,
    ansi_color(GREEN, BLACK, strings.Join(post_values, "&")),
    ansi_color(YELLO, BLACK, query_url),
    ansi_color(YELLO, BLACK, strings.Join(post_values, "&")),
  )
  g_env.AccessLog.Println(log_line)
}

//transit complete logger
func access_log(aid string, w http.ResponseWriter, r *http.Request, query_url string, startTime time.Time) {
  post_values := post_values(r)

  if len(post_values) == 0 {
    post_values = append(post_values, "-")
  }

  content_len := w.Header().Get("Content-Length")
  if len(content_len) == 0 {
    content_len = "-"
  }

  log_line := fmt.Sprintf(`Complete [%s] "%s" [%s] S:"%s %s %s F:{%s}" D:"%s F:{%s} %s %s",%0.5fs`,
    ansi_color(WHITE, BLACK, access_ip(r)),
    ansi_color(WHITE, RED, aid),
    time.Now().Format("2006-01-02 15:04:05.999999999 -0700 MST"),
    r.Method,
    ansi_color(GREEN, BLACK, r.RequestURI),
    r.Proto,
    ansi_color(GREEN, BLACK, strings.Join(post_values, "&")),
    ansi_color(YELLO, BLACK, query_url),
    ansi_color(YELLO, BLACK, strings.Join(post_values, "&")),
    w.Header().Get("Status"),
    content_len,
    time.Now().Sub(startTime).Seconds(),
  )
  g_env.AccessLog.Println(log_line)
}

func parse_querys(r *http.Request) (raw_query []string) {
  r.ParseForm()
  check_key_map := make(map[string]bool)
  check_and_append := func(_k, _v string) {
    if _, ok := check_key_map[_k]; ok {
      return
    }
    check_key_map[_k] = true
    raw_query = append(raw_query, fmt.Sprintf("%s=%s", url.QueryEscape(_k), url.QueryEscape(_v)))
  }

  for k, _ := range r.Form {
    check_and_append(k, r.Form.Get(k))
  }

  if len(r.Referer()) > 0 {
    if uri, err := url.Parse(r.Referer()); err == nil {
      for k, _ := range uri.Query() {
        check_and_append(k, uri.Query().Get(k))
      }
    }
  }
  return
}

/**
获取目标地址
*/
func target_path(r *http.Request, cfg *config.Config) (t string) {
  if len(cfg.TargetPath) > 0 {
    t = cfg.TargetPath
  } else {
    t = r.URL.Path
  }
  return
}

/**
获取目标服务服务器
*/
func target_server(cfg *config.Config) (s string) {
  if len(cfg.TargetServer) > 0 {
    s = cfg.TargetServer
  } else {
    s = g_config.Default.TargetServer
  }
  return
}

/*
标准http解析库实现
*/
func swap_raw_query(r *http.Request, cfg *config.Config) (q string) {
  var tmp_slice []string
  append_slict := func(key string, value string) {
    tmp_slice = append(tmp_slice, fmt.Sprintf("%s=%s", url.QueryEscape(key), url.QueryEscape(r.URL.Query().Get(value))))
  }
  //如果配置段没有需要交换的参数,则直接返回查询字符串
  if len(cfg.TargetParamNameSwap) == 0 {
    for k, _ := range r.URL.Query() {
      append_slict(k, k)
    }
    q = strings.Join(tmp_slice, "&")
    return
  }

  for k, _ := range r.URL.Query() {
    if v, ok := cfg.TargetParamNameSwap[k]; ok {
      append_slict(v, k)
    } else {
      append_slict(k, k)
    }
  }
  q = strings.Join(tmp_slice, "&")
  return
}

func timeout_dialer(conn_timeout int, rw_timeout int) func(net, addr string) (c net.Conn, err error) {
  return func(netw, addr string) (net.Conn, error) {
    conn, err := net.DialTimeout(netw, addr, time.Duration(conn_timeout)*time.Second)
    if err != nil {
      log.Printf("Failed to connect to [%s]. Timed out after %d seconds\n", addr, rw_timeout)
      return nil, err
    }
    conn.SetDeadline(time.Now().Add(time.Duration(rw_timeout) * time.Second))
    return conn, nil
  }
}

func (s Server) handler_func(w http.ResponseWriter, r *http.Request) {
  var (
    cfg                                  *config.Config
    cfg_err                              *config.ConfigErr
    conntction_timeout, response_timeout int
    req                                  *http.Request
    err                                  error
    raw_query                            []string
  )
  defer func() {
    if re := recover(); re != nil {
      g_env.ErrorLog.Println("Recovered in backendServer:", re)
    }
  }()

  defer r.Body.Close()

  start_at := time.Now()
  raw_query = parse_querys(r)
  aid := access_id(r.RequestURI)

  if err != nil {
    g_env.ErrorLog.Println(req, err)
    http.Error(w, "Read Body Error.", http.StatusInternalServerError)
    return
  }

  //获取配置文件
  if cfg, cfg_err = g_config.FindBySourcePathAndParams(raw_query, r.URL.Path); cfg_err != nil {
    cfg = g_config.FindByParamsOrSourcePath(raw_query, r.URL.Path)
  }

  if conntction_timeout = cfg.ConnectionTimeout; conntction_timeout <= 0 {
    conntction_timeout = 15
  }

  if response_timeout = cfg.ResponseTimeout; response_timeout <= 0 {
    response_timeout = 120
  }

  transport := http.Transport{
    Dial: timeout_dialer(conntction_timeout, response_timeout),
    ResponseHeaderTimeout: time.Duration(response_timeout) * time.Second,
    DisableCompression:    false,
    DisableKeepAlives:     true,
    MaxIdleConnsPerHost:   2,
  }
  defer transport.CloseIdleConnections()

  client := &http.Client{
    Transport: &transport,
  }

  query_url, _ := url.Parse(target_server(cfg) + target_path(r, cfg) + "?" + swap_raw_query(r, cfg))

  access_log_begin(aid, r, query_url.String())

  switch r.Method {
  case "GET", "HEAD":
    req, err = http.NewRequest(r.Method, query_url.String(), nil)
  case "POST":
    req, err = http.NewRequest(r.Method, query_url.String(), bytes.NewBufferString(strings.Join(raw_query, "&")))
    req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
  default:
    http.Error(w, "MethodNotAllowed", http.StatusMethodNotAllowed)
    return
  }
  req.Close = true

  header_copy(r.Header, &req.Header)

  if err != nil {
    g_env.ErrorLog.Println(err)
    http.Error(w, err.Error(), http.StatusInternalServerError)
    return
  }

  resp, err := client.Do(req)
  defer resp.Body.Close()

  if err != nil {
    g_env.ErrorLog.Println(req, err)
    http.Error(w, err.Error(), http.StatusInternalServerError)
    return
  }

  for hk, _ := range resp.Header {
    w.Header().Set(hk, resp.Header.Get(hk))
  }
  w.Header().Set("X-Transit-Ver", "0.0.1")
  w.Header().Set("Server", "X-Transit")

  w.WriteHeader(resp.StatusCode)
  io.Copy(w, resp.Body)
  access_log(aid, w, r, query_url.String(), start_at)
}

//TODO
// add unix socket listener

func Run() {
  g_env.ErrorLog.Printf("start@ %s:%d %v \n", g_config.Listen.Host, g_config.Listen.Port, time.Now())
  fmt.Printf("start@ %s:%d %v \n", g_config.Listen.Host, g_config.Listen.Port, time.Now())

  sigchan := make(chan os.Signal, 1)
  signal.Notify(sigchan, os.Interrupt, syscall.SIGTERM, syscall.SIGHUP, syscall.SIGINT, syscall.SIGQUIT)

  server := Server{}
  go func() {
    s := &http.Server{
      Addr:           fmt.Sprintf("%s:%d", g_config.Listen.Host, g_config.Listen.Port),
      Handler:        http.HandlerFunc(server.handler_func),
      ReadTimeout:    120 * time.Second,
      WriteTimeout:   120 * time.Second,
      MaxHeaderBytes: 1 << 20,
    }
    log.Fatal(s.ListenAndServe())
  }()

  <-sigchan
}
<|endoftext|>


<commit_before>/*
 * Copyright (C) Max Romanov
 * Copyright (C) NGINX, Inc.
 */

package unit

/*
#include "nxt_cgo_lib.h"
*/
import "C"

import (
	"net/http"
)

type response struct {
	header     http.Header
	headerSent bool
	req        *http.Request
	c_req      C.uintptr_t
}

func new_response(c_req C.uintptr_t, req *http.Request) *response {
	resp := &response{
		header: http.Header{},
		req:    req,
		c_req:  c_req,
	}

	return resp
}

func (r *response) Header() http.Header {
	return r.header
}

func (r *response) Write(p []byte) (n int, err error) {
	if !r.headerSent {
		r.WriteHeader(http.StatusOK)
	}

	res := C.nxt_cgo_response_write(r.c_req, buf_ref(p), C.uint32_t(len(p)))
	return int(res), nil
}

func (r *response) WriteHeader(code int) {
	if r.headerSent {
		// Note: explicitly using Stderr, as Stdout is our HTTP output.
		nxt_go_warn("multiple response.WriteHeader calls")
		return
	}
	r.headerSent = true

	// Set a default Content-Type
	if _, hasType := r.header["Content-Type"]; !hasType {
		r.header.Add("Content-Type", "text/html; charset=utf-8")
	}

	fields := 0
	fields_size := 0

	for k, vv := range r.header {
		for _, v := range vv {
			fields++
			fields_size += len(k) + len(v)
		}
	}

	C.nxt_cgo_response_create(r.c_req, C.int(code), C.int(fields),
		C.uint32_t(fields_size))

	for k, vv := range r.header {
		for _, v := range vv {
			C.nxt_cgo_response_add_field(r.c_req, str_ref(k), C.uint8_t(len(k)),
				str_ref(v), C.uint32_t(len(v)))
		}
	}

	C.nxt_cgo_response_send(r.c_req)
}

func (r *response) Flush() {
	if !r.headerSent {
		r.WriteHeader(http.StatusOK)
	}
}
<commit_msg>Go: fixing header buffer size calculation.<commit_after>/*
 * Copyright (C) Max Romanov
 * Copyright (C) NGINX, Inc.
 */

package unit

/*
#include "nxt_cgo_lib.h"
*/
import "C"

import (
	"net/http"
)

type response struct {
	header     http.Header
	headerSent bool
	req        *http.Request
	c_req      C.uintptr_t
}

func new_response(c_req C.uintptr_t, req *http.Request) *response {
	resp := &response{
		header: http.Header{},
		req:    req,
		c_req:  c_req,
	}

	return resp
}

func (r *response) Header() http.Header {
	return r.header
}

func (r *response) Write(p []byte) (n int, err error) {
	if !r.headerSent {
		r.WriteHeader(http.StatusOK)
	}

	res := C.nxt_cgo_response_write(r.c_req, buf_ref(p), C.uint32_t(len(p)))
	return int(res), nil
}

func (r *response) WriteHeader(code int) {
	if r.headerSent {
		// Note: explicitly using Stderr, as Stdout is our HTTP output.
		nxt_go_warn("multiple response.WriteHeader calls")
		return
	}
	r.headerSent = true

	// Set a default Content-Type
	if _, hasType := r.header["Content-Type"]; !hasType {
		r.header.Add("Content-Type", "text/html; charset=utf-8")
	}

	fields := 0
	fields_size := 0

	for k, vv := range r.header {
		for _, v := range vv {
			fields++
			fields_size += len(k) + len(v) + 2
		}
	}

	C.nxt_cgo_response_create(r.c_req, C.int(code), C.int(fields),
		C.uint32_t(fields_size))

	for k, vv := range r.header {
		for _, v := range vv {
			C.nxt_cgo_response_add_field(r.c_req, str_ref(k), C.uint8_t(len(k)),
				str_ref(v), C.uint32_t(len(v)))
		}
	}

	C.nxt_cgo_response_send(r.c_req)
}

func (r *response) Flush() {
	if !r.headerSent {
		r.WriteHeader(http.StatusOK)
	}
}
<|endoftext|>


<commit_before>/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you 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 thrift

import (
	"log"
	"runtime/debug"
	"sync"
)

/*
 * This is not a typical TSimpleServer as it is not blocked after accept a socket.
 * It is more like a TThreadedServer that can handle different connections in different goroutines.
 * This will work if golang user implements a conn-pool like thing in client side.
 */
type TSimpleServer struct {
	quit chan struct{}

	processorFactory       TProcessorFactory
	serverTransport        TServerTransport
	inputTransportFactory  TTransportFactory
	outputTransportFactory TTransportFactory
	inputProtocolFactory   TProtocolFactory
	outputProtocolFactory  TProtocolFactory
}

func NewTSimpleServer2(processor TProcessor, serverTransport TServerTransport) *TSimpleServer {
	return NewTSimpleServerFactory2(NewTProcessorFactory(processor), serverTransport)
}

func NewTSimpleServer4(processor TProcessor, serverTransport TServerTransport, transportFactory TTransportFactory, protocolFactory TProtocolFactory) *TSimpleServer {
	return NewTSimpleServerFactory4(NewTProcessorFactory(processor),
		serverTransport,
		transportFactory,
		protocolFactory,
	)
}

func NewTSimpleServer6(processor TProcessor, serverTransport TServerTransport, inputTransportFactory TTransportFactory, outputTransportFactory TTransportFactory, inputProtocolFactory TProtocolFactory, outputProtocolFactory TProtocolFactory) *TSimpleServer {
	return NewTSimpleServerFactory6(NewTProcessorFactory(processor),
		serverTransport,
		inputTransportFactory,
		outputTransportFactory,
		inputProtocolFactory,
		outputProtocolFactory,
	)
}

func NewTSimpleServerFactory2(processorFactory TProcessorFactory, serverTransport TServerTransport) *TSimpleServer {
	return NewTSimpleServerFactory6(processorFactory,
		serverTransport,
		NewTTransportFactory(),
		NewTTransportFactory(),
		NewTBinaryProtocolFactoryDefault(),
		NewTBinaryProtocolFactoryDefault(),
	)
}

func NewTSimpleServerFactory4(processorFactory TProcessorFactory, serverTransport TServerTransport, transportFactory TTransportFactory, protocolFactory TProtocolFactory) *TSimpleServer {
	return NewTSimpleServerFactory6(processorFactory,
		serverTransport,
		transportFactory,
		transportFactory,
		protocolFactory,
		protocolFactory,
	)
}

func NewTSimpleServerFactory6(processorFactory TProcessorFactory, serverTransport TServerTransport, inputTransportFactory TTransportFactory, outputTransportFactory TTransportFactory, inputProtocolFactory TProtocolFactory, outputProtocolFactory TProtocolFactory) *TSimpleServer {
	return &TSimpleServer{
		processorFactory:       processorFactory,
		serverTransport:        serverTransport,
		inputTransportFactory:  inputTransportFactory,
		outputTransportFactory: outputTransportFactory,
		inputProtocolFactory:   inputProtocolFactory,
		outputProtocolFactory:  outputProtocolFactory,
		quit: make(chan struct{}, 1),
	}
}

func (p *TSimpleServer) ProcessorFactory() TProcessorFactory {
	return p.processorFactory
}

func (p *TSimpleServer) ServerTransport() TServerTransport {
	return p.serverTransport
}

func (p *TSimpleServer) InputTransportFactory() TTransportFactory {
	return p.inputTransportFactory
}

func (p *TSimpleServer) OutputTransportFactory() TTransportFactory {
	return p.outputTransportFactory
}

func (p *TSimpleServer) InputProtocolFactory() TProtocolFactory {
	return p.inputProtocolFactory
}

func (p *TSimpleServer) OutputProtocolFactory() TProtocolFactory {
	return p.outputProtocolFactory
}

func (p *TSimpleServer) Listen() error {
	return p.serverTransport.Listen()
}

func (p *TSimpleServer) AcceptLoop() error {
	for {
		client, err := p.serverTransport.Accept()
		if err != nil {
			select {
			case <-p.quit:
				return nil
			default:
			}
			return err
		}
		if client != nil {
			go func() {
				if err := p.processRequests(client); err != nil {
					log.Println("error processing request:", err)
				}
			}()
		}
	}
}

func (p *TSimpleServer) Serve() error {
	err := p.Listen()
	if err != nil {
		return err
	}
	p.AcceptLoop()
	return nil
}

var once sync.Once

func (p *TSimpleServer) Stop() error {
	q := func() {
		p.quit <- struct{}{}
		p.serverTransport.Interrupt()
	}
	once.Do(q)
	return nil
}

func (p *TSimpleServer) processRequests(client TTransport) error {
	processor := p.processorFactory.GetProcessor(client)
	inputTransport := p.inputTransportFactory.GetTransport(client)
	outputTransport := p.outputTransportFactory.GetTransport(client)
	inputProtocol := p.inputProtocolFactory.GetProtocol(inputTransport)
	outputProtocol := p.outputProtocolFactory.GetProtocol(outputTransport)
	defer func() {
		if e := recover(); e != nil {
			log.Printf("panic in processor: %s: %s", e, debug.Stack())
		}
	}()
	if inputTransport != nil {
		defer inputTransport.Close()
	}
	if outputTransport != nil {
		defer outputTransport.Close()
	}
	for {
		ok, err := processor.Process(inputProtocol, outputProtocol)
		if err, ok := err.(TTransportException); ok && err.TypeId() == END_OF_FILE {
			return nil
		} else if err != nil {
			log.Printf("error processing request: %s", err)
			return err
		}
		if err, ok := err.(TApplicationException); ok && err.TypeId() == UNKNOWN_METHOD {
			continue
		}
 		if !ok {
			break
		}
	}
	return nil
}
<commit_msg>THRIFT-3730: go server errors are logged twice Client: go Patch: zhujun2006 <stone1342006 at gmail dot com><commit_after>/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. The ASF licenses this file
 * to you 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 thrift

import (
	"log"
	"runtime/debug"
	"sync"
)

/*
 * This is not a typical TSimpleServer as it is not blocked after accept a socket.
 * It is more like a TThreadedServer that can handle different connections in different goroutines.
 * This will work if golang user implements a conn-pool like thing in client side.
 */
type TSimpleServer struct {
	quit chan struct{}

	processorFactory       TProcessorFactory
	serverTransport        TServerTransport
	inputTransportFactory  TTransportFactory
	outputTransportFactory TTransportFactory
	inputProtocolFactory   TProtocolFactory
	outputProtocolFactory  TProtocolFactory
}

func NewTSimpleServer2(processor TProcessor, serverTransport TServerTransport) *TSimpleServer {
	return NewTSimpleServerFactory2(NewTProcessorFactory(processor), serverTransport)
}

func NewTSimpleServer4(processor TProcessor, serverTransport TServerTransport, transportFactory TTransportFactory, protocolFactory TProtocolFactory) *TSimpleServer {
	return NewTSimpleServerFactory4(NewTProcessorFactory(processor),
		serverTransport,
		transportFactory,
		protocolFactory,
	)
}

func NewTSimpleServer6(processor TProcessor, serverTransport TServerTransport, inputTransportFactory TTransportFactory, outputTransportFactory TTransportFactory, inputProtocolFactory TProtocolFactory, outputProtocolFactory TProtocolFactory) *TSimpleServer {
	return NewTSimpleServerFactory6(NewTProcessorFactory(processor),
		serverTransport,
		inputTransportFactory,
		outputTransportFactory,
		inputProtocolFactory,
		outputProtocolFactory,
	)
}

func NewTSimpleServerFactory2(processorFactory TProcessorFactory, serverTransport TServerTransport) *TSimpleServer {
	return NewTSimpleServerFactory6(processorFactory,
		serverTransport,
		NewTTransportFactory(),
		NewTTransportFactory(),
		NewTBinaryProtocolFactoryDefault(),
		NewTBinaryProtocolFactoryDefault(),
	)
}

func NewTSimpleServerFactory4(processorFactory TProcessorFactory, serverTransport TServerTransport, transportFactory TTransportFactory, protocolFactory TProtocolFactory) *TSimpleServer {
	return NewTSimpleServerFactory6(processorFactory,
		serverTransport,
		transportFactory,
		transportFactory,
		protocolFactory,
		protocolFactory,
	)
}

func NewTSimpleServerFactory6(processorFactory TProcessorFactory, serverTransport TServerTransport, inputTransportFactory TTransportFactory, outputTransportFactory TTransportFactory, inputProtocolFactory TProtocolFactory, outputProtocolFactory TProtocolFactory) *TSimpleServer {
	return &TSimpleServer{
		processorFactory:       processorFactory,
		serverTransport:        serverTransport,
		inputTransportFactory:  inputTransportFactory,
		outputTransportFactory: outputTransportFactory,
		inputProtocolFactory:   inputProtocolFactory,
		outputProtocolFactory:  outputProtocolFactory,
		quit: make(chan struct{}, 1),
	}
}

func (p *TSimpleServer) ProcessorFactory() TProcessorFactory {
	return p.processorFactory
}

func (p *TSimpleServer) ServerTransport() TServerTransport {
	return p.serverTransport
}

func (p *TSimpleServer) InputTransportFactory() TTransportFactory {
	return p.inputTransportFactory
}

func (p *TSimpleServer) OutputTransportFactory() TTransportFactory {
	return p.outputTransportFactory
}

func (p *TSimpleServer) InputProtocolFactory() TProtocolFactory {
	return p.inputProtocolFactory
}

func (p *TSimpleServer) OutputProtocolFactory() TProtocolFactory {
	return p.outputProtocolFactory
}

func (p *TSimpleServer) Listen() error {
	return p.serverTransport.Listen()
}

func (p *TSimpleServer) AcceptLoop() error {
	for {
		client, err := p.serverTransport.Accept()
		if err != nil {
			select {
			case <-p.quit:
				return nil
			default:
			}
			return err
		}
		if client != nil {
			go func() {
				if err := p.processRequests(client); err != nil {
					log.Println("error processing request:", err)
				}
			}()
		}
	}
}

func (p *TSimpleServer) Serve() error {
	err := p.Listen()
	if err != nil {
		return err
	}
	p.AcceptLoop()
	return nil
}

var once sync.Once

func (p *TSimpleServer) Stop() error {
	q := func() {
		p.quit <- struct{}{}
		p.serverTransport.Interrupt()
	}
	once.Do(q)
	return nil
}

func (p *TSimpleServer) processRequests(client TTransport) error {
	processor := p.processorFactory.GetProcessor(client)
	inputTransport := p.inputTransportFactory.GetTransport(client)
	outputTransport := p.outputTransportFactory.GetTransport(client)
	inputProtocol := p.inputProtocolFactory.GetProtocol(inputTransport)
	outputProtocol := p.outputProtocolFactory.GetProtocol(outputTransport)
	defer func() {
		if e := recover(); e != nil {
			log.Printf("panic in processor: %s: %s", e, debug.Stack())
		}
	}()
	if inputTransport != nil {
		defer inputTransport.Close()
	}
	if outputTransport != nil {
		defer outputTransport.Close()
	}
	for {
		ok, err := processor.Process(inputProtocol, outputProtocol)
		if err, ok := err.(TTransportException); ok && err.TypeId() == END_OF_FILE {
			return nil
		} else if err != nil {
			return err
		}
		if err, ok := err.(TApplicationException); ok && err.TypeId() == UNKNOWN_METHOD {
			continue
		}
 		if !ok {
			break
		}
	}
	return nil
}
<|endoftext|>


<commit_before>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2020 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package types

import (
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestHostnameTrieInsert(t *testing.T) {
	hostnames := HostnameTrie{}
	assert.NoError(t, hostnames.insert("test.k6.io"))
	assert.Error(t, hostnames.insert("inval*d.pattern"))
	assert.NoError(t, hostnames.insert("*valid.pattern"))
}

func TestHostnameTrieContains(t *testing.T) {
	trie, err := NewHostnameTrie([]string{"test.k6.io", "*valid.pattern"})
	require.NoError(t, err)
	_, matches := trie.Contains("K6.Io")
	assert.False(t, matches)
	match, matches := trie.Contains("tEsT.k6.Io")
	assert.True(t, matches)
	assert.Equal(t, "test.k6.io", match)
	match, matches = trie.Contains("TEST.K6.IO")
	assert.True(t, matches)
	assert.Equal(t, "test.k6.io", match)
	match, matches = trie.Contains("blocked.valId.paTtern")
	assert.True(t, matches)
	assert.Equal(t, "*valid.pattern", match)
	_, matches = trie.Contains("example.test.k6.io")
	assert.False(t, matches)
}
<commit_msg>rewrite TestHostnameTrieContains as table test<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2020 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package types

import (
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestHostnameTrieInsert(t *testing.T) {
	hostnames := HostnameTrie{}
	assert.NoError(t, hostnames.insert("test.k6.io"))
	assert.Error(t, hostnames.insert("inval*d.pattern"))
	assert.NoError(t, hostnames.insert("*valid.pattern"))
}

func TestHostnameTrieContains(t *testing.T) {
	trie, err := NewHostnameTrie([]string{"test.k6.io", "*valid.pattern"})
	require.NoError(t, err)
	cases := map[string]string{
		"K6.Io":                 "",
		"tEsT.k6.Io":            "test.k6.io",
		"TESt.K6.IO":            "test.k6.io",
		"blocked.valId.paTtern": "*valid.pattern",
		"example.test.k6.io":    "",
	}
	for key, value := range cases {
		host, pattern := key, value
		t.Run(host, func(t *testing.T) {
			match, matches := trie.Contains(host)
			if pattern == "" {
				assert.False(t, matches)
				assert.Empty(t, match)
			} else {
				assert.True(t, matches)
				assert.Equal(t, pattern, match)
			}
		})
	}
}
<|endoftext|>


<commit_before>package configs

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os/exec"
	"time"

	"github.com/Sirupsen/logrus"
)

type Rlimit struct {
	Type int    `json:"type"`
	Hard uint64 `json:"hard"`
	Soft uint64 `json:"soft"`
}

// IDMap represents UID/GID Mappings for User Namespaces.
type IDMap struct {
	ContainerID int `json:"container_id"`
	HostID      int `json:"host_id"`
	Size        int `json:"size"`
}

// Seccomp represents syscall restrictions
// By default, only the native architecture of the kernel is allowed to be used
// for syscalls. Additional architectures can be added by specifying them in
// Architectures.
type Seccomp struct {
	DefaultAction Action     `json:"default_action"`
	Architectures []string   `json:"architectures"`
	Syscalls      []*Syscall `json:"syscalls"`
}

// An action to be taken upon rule match in Seccomp
type Action int

const (
	Kill Action = iota + 1
	Errno
	Trap
	Allow
	Trace
)

// A comparison operator to be used when matching syscall arguments in Seccomp
type Operator int

const (
	EqualTo Operator = iota + 1
	NotEqualTo
	GreaterThan
	GreaterThanOrEqualTo
	LessThan
	LessThanOrEqualTo
	MaskEqualTo
)

// A rule to match a specific syscall argument in Seccomp
type Arg struct {
	Index    uint     `json:"index"`
	Value    uint64   `json:"value"`
	ValueTwo uint64   `json:"value_two"`
	Op       Operator `json:"op"`
}

// An rule to match a syscall in Seccomp
type Syscall struct {
	Name   string `json:"name"`
	Action Action `json:"action"`
	Args   []*Arg `json:"args"`
}

// TODO Windows. Many of these fields should be factored out into those parts
// which are common across platforms, and those which are platform specific.

// Config defines configuration options for executing a process inside a contained environment.
type Config struct {
	// NoPivotRoot will use MS_MOVE and a chroot to jail the process into the container's rootfs
	// This is a common option when the container is running in ramdisk
	NoPivotRoot bool `json:"no_pivot_root"`

	// ParentDeathSignal specifies the signal that is sent to the container's process in the case
	// that the parent process dies.
	ParentDeathSignal int `json:"parent_death_signal"`

	// PivotDir allows a custom directory inside the container's root filesystem to be used as pivot, when NoPivotRoot is not set.
	// When a custom PivotDir not set, a temporary dir inside the root filesystem will be used. The pivot dir needs to be writeable.
	// This is required when using read only root filesystems. In these cases, a read/writeable path can be (bind) mounted somewhere inside the root filesystem to act as pivot.
	PivotDir string `json:"pivot_dir"`

	// Path to a directory containing the container's root filesystem.
	Rootfs string `json:"rootfs"`

	// Readonlyfs will remount the container's rootfs as readonly where only externally mounted
	// bind mounts are writtable.
	Readonlyfs bool `json:"readonlyfs"`

	// Specifies the mount propagation flags to be applied to /.
	RootPropagation int `json:"rootPropagation"`

	// Mounts specify additional source and destination paths that will be mounted inside the container's
	// rootfs and mount namespace if specified
	Mounts []*Mount `json:"mounts"`

	// The device nodes that should be automatically created within the container upon container start.  Note, make sure that the node is marked as allowed in the cgroup as well!
	Devices []*Device `json:"devices"`

	MountLabel string `json:"mount_label"`

	// Hostname optionally sets the container's hostname if provided
	Hostname string `json:"hostname"`

	// Namespaces specifies the container's namespaces that it should setup when cloning the init process
	// If a namespace is not provided that namespace is shared from the container's parent process
	Namespaces Namespaces `json:"namespaces"`

	// Capabilities specify the capabilities to keep when executing the process inside the container
	// All capbilities not specified will be dropped from the processes capability mask
	Capabilities []string `json:"capabilities"`

	// Networks specifies the container's network setup to be created
	Networks []*Network `json:"networks"`

	// Routes can be specified to create entries in the route table as the container is started
	Routes []*Route `json:"routes"`

	// Cgroups specifies specific cgroup settings for the various subsystems that the container is
	// placed into to limit the resources the container has available
	Cgroups *Cgroup `json:"cgroups"`

	// AppArmorProfile specifies the profile to apply to the process running in the container and is
	// change at the time the process is execed
	AppArmorProfile string `json:"apparmor_profile,omitempty"`

	// ProcessLabel specifies the label to apply to the process running in the container.  It is
	// commonly used by selinux
	ProcessLabel string `json:"process_label,omitempty"`

	// Rlimits specifies the resource limits, such as max open files, to set in the container
	// If Rlimits are not set, the container will inherit rlimits from the parent process
	Rlimits []Rlimit `json:"rlimits,omitempty"`

	// OomScoreAdj specifies the adjustment to be made by the kernel when calculating oom scores
	// for a process. Valid values are between the range [-1000, '1000'], where processes with
	// higher scores are preferred for being killed.
	// More information about kernel oom score calculation here: https://lwn.net/Articles/317814/
	OomScoreAdj int `json:"oom_score_adj"`

	// AdditionalGroups specifies the gids that should be added to supplementary groups
	// in addition to those that the user belongs to.
	AdditionalGroups []string `json:"additional_groups"`

	// UidMappings is an array of User ID mappings for User Namespaces
	UidMappings []IDMap `json:"uid_mappings"`

	// GidMappings is an array of Group ID mappings for User Namespaces
	GidMappings []IDMap `json:"gid_mappings"`

	// MaskPaths specifies paths within the container's rootfs to mask over with a bind
	// mount pointing to /dev/null as to prevent reads of the file.
	MaskPaths []string `json:"mask_paths"`

	// ReadonlyPaths specifies paths within the container's rootfs to remount as read-only
	// so that these files prevent any writes.
	ReadonlyPaths []string `json:"readonly_paths"`

	// Sysctl is a map of properties and their values. It is the equivalent of using
	// sysctl -w my.property.name value in Linux.
	Sysctl map[string]string `json:"sysctl"`

	// Seccomp allows actions to be taken whenever a syscall is made within the container.
	// A number of rules are given, each having an action to be taken if a syscall matches it.
	// A default action to be taken if no rules match is also given.
	Seccomp *Seccomp `json:"seccomp"`

	// NoNewPrivileges controls whether processes in the container can gain additional privileges.
	NoNewPrivileges bool `json:"no_new_privileges,omitempty"`

	// Hooks are a collection of actions to perform at various container lifecycle events.
	// CommandHooks are serialized to JSON, but other hooks are not.
	Hooks *Hooks

	// Version is the version of opencontainer specification that is supported.
	Version string `json:"version"`

	// Labels are user defined metadata that is stored in the config and populated on the state
	Labels []string `json:"labels"`
}

type Hooks struct {
	// Prestart commands are executed after the container namespaces are created,
	// but before the user supplied command is executed from init.
	Prestart []Hook

	// Poststart commands are executed after the container init process starts.
	Poststart []Hook

	// Poststop commands are executed after the container init process exits.
	Poststop []Hook
}

func (hooks *Hooks) UnmarshalJSON(b []byte) error {
	var state struct {
		Prestart  []CommandHook
		Poststart []CommandHook
		Poststop  []CommandHook
	}

	if err := json.Unmarshal(b, &state); err != nil {
		return err
	}

	deserialize := func(shooks []CommandHook) (hooks []Hook) {
		for _, shook := range shooks {
			hooks = append(hooks, shook)
		}

		return hooks
	}

	hooks.Prestart = deserialize(state.Prestart)
	hooks.Poststart = deserialize(state.Poststart)
	hooks.Poststop = deserialize(state.Poststop)
	return nil
}

func (hooks Hooks) MarshalJSON() ([]byte, error) {
	serialize := func(hooks []Hook) (serializableHooks []CommandHook) {
		for _, hook := range hooks {
			switch chook := hook.(type) {
			case CommandHook:
				serializableHooks = append(serializableHooks, chook)
			default:
				logrus.Warnf("cannot serialize hook of type %T, skipping", hook)
			}
		}

		return serializableHooks
	}

	return json.Marshal(map[string]interface{}{
		"prestart":  serialize(hooks.Prestart),
		"poststart": serialize(hooks.Poststart),
		"poststop":  serialize(hooks.Poststop),
	})
}

// HookState is the payload provided to a hook on execution.
type HookState struct {
	Version string `json:"version"`
	ID      string `json:"id"`
	Pid     int    `json:"pid"`
	Root    string `json:"root"`
}

type Hook interface {
	// Run executes the hook with the provided state.
	Run(HookState) error
}

// NewFunctionHooks will call the provided function when the hook is run.
func NewFunctionHook(f func(HookState) error) FuncHook {
	return FuncHook{
		run: f,
	}
}

type FuncHook struct {
	run func(HookState) error
}

func (f FuncHook) Run(s HookState) error {
	return f.run(s)
}

type Command struct {
	Path    string         `json:"path"`
	Args    []string       `json:"args"`
	Env     []string       `json:"env"`
	Dir     string         `json:"dir"`
	Timeout *time.Duration `json:"timeout"`
}

// NewCommandHooks will execute the provided command when the hook is run.
func NewCommandHook(cmd Command) CommandHook {
	return CommandHook{
		Command: cmd,
	}
}

type CommandHook struct {
	Command
}

func (c Command) Run(s HookState) error {
	b, err := json.Marshal(s)
	if err != nil {
		return err
	}
	cmd := exec.Cmd{
		Path:  c.Path,
		Args:  c.Args,
		Env:   c.Env,
		Stdin: bytes.NewReader(b),
	}
	errC := make(chan error, 1)
	go func() {
		errC <- cmd.Run()
	}()
	if c.Timeout != nil {
		select {
		case err := <-errC:
			return err
		case <-time.After(*c.Timeout):
			cmd.Process.Kill()
			cmd.Wait()
			return fmt.Errorf("hook ran past specified timeout of %.1fs", c.Timeout.Seconds())
		}
	}
	return <-errC
}
<commit_msg>Report hook output on error<commit_after>package configs

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os/exec"
	"time"

	"github.com/Sirupsen/logrus"
)

type Rlimit struct {
	Type int    `json:"type"`
	Hard uint64 `json:"hard"`
	Soft uint64 `json:"soft"`
}

// IDMap represents UID/GID Mappings for User Namespaces.
type IDMap struct {
	ContainerID int `json:"container_id"`
	HostID      int `json:"host_id"`
	Size        int `json:"size"`
}

// Seccomp represents syscall restrictions
// By default, only the native architecture of the kernel is allowed to be used
// for syscalls. Additional architectures can be added by specifying them in
// Architectures.
type Seccomp struct {
	DefaultAction Action     `json:"default_action"`
	Architectures []string   `json:"architectures"`
	Syscalls      []*Syscall `json:"syscalls"`
}

// An action to be taken upon rule match in Seccomp
type Action int

const (
	Kill Action = iota + 1
	Errno
	Trap
	Allow
	Trace
)

// A comparison operator to be used when matching syscall arguments in Seccomp
type Operator int

const (
	EqualTo Operator = iota + 1
	NotEqualTo
	GreaterThan
	GreaterThanOrEqualTo
	LessThan
	LessThanOrEqualTo
	MaskEqualTo
)

// A rule to match a specific syscall argument in Seccomp
type Arg struct {
	Index    uint     `json:"index"`
	Value    uint64   `json:"value"`
	ValueTwo uint64   `json:"value_two"`
	Op       Operator `json:"op"`
}

// An rule to match a syscall in Seccomp
type Syscall struct {
	Name   string `json:"name"`
	Action Action `json:"action"`
	Args   []*Arg `json:"args"`
}

// TODO Windows. Many of these fields should be factored out into those parts
// which are common across platforms, and those which are platform specific.

// Config defines configuration options for executing a process inside a contained environment.
type Config struct {
	// NoPivotRoot will use MS_MOVE and a chroot to jail the process into the container's rootfs
	// This is a common option when the container is running in ramdisk
	NoPivotRoot bool `json:"no_pivot_root"`

	// ParentDeathSignal specifies the signal that is sent to the container's process in the case
	// that the parent process dies.
	ParentDeathSignal int `json:"parent_death_signal"`

	// PivotDir allows a custom directory inside the container's root filesystem to be used as pivot, when NoPivotRoot is not set.
	// When a custom PivotDir not set, a temporary dir inside the root filesystem will be used. The pivot dir needs to be writeable.
	// This is required when using read only root filesystems. In these cases, a read/writeable path can be (bind) mounted somewhere inside the root filesystem to act as pivot.
	PivotDir string `json:"pivot_dir"`

	// Path to a directory containing the container's root filesystem.
	Rootfs string `json:"rootfs"`

	// Readonlyfs will remount the container's rootfs as readonly where only externally mounted
	// bind mounts are writtable.
	Readonlyfs bool `json:"readonlyfs"`

	// Specifies the mount propagation flags to be applied to /.
	RootPropagation int `json:"rootPropagation"`

	// Mounts specify additional source and destination paths that will be mounted inside the container's
	// rootfs and mount namespace if specified
	Mounts []*Mount `json:"mounts"`

	// The device nodes that should be automatically created within the container upon container start.  Note, make sure that the node is marked as allowed in the cgroup as well!
	Devices []*Device `json:"devices"`

	MountLabel string `json:"mount_label"`

	// Hostname optionally sets the container's hostname if provided
	Hostname string `json:"hostname"`

	// Namespaces specifies the container's namespaces that it should setup when cloning the init process
	// If a namespace is not provided that namespace is shared from the container's parent process
	Namespaces Namespaces `json:"namespaces"`

	// Capabilities specify the capabilities to keep when executing the process inside the container
	// All capbilities not specified will be dropped from the processes capability mask
	Capabilities []string `json:"capabilities"`

	// Networks specifies the container's network setup to be created
	Networks []*Network `json:"networks"`

	// Routes can be specified to create entries in the route table as the container is started
	Routes []*Route `json:"routes"`

	// Cgroups specifies specific cgroup settings for the various subsystems that the container is
	// placed into to limit the resources the container has available
	Cgroups *Cgroup `json:"cgroups"`

	// AppArmorProfile specifies the profile to apply to the process running in the container and is
	// change at the time the process is execed
	AppArmorProfile string `json:"apparmor_profile,omitempty"`

	// ProcessLabel specifies the label to apply to the process running in the container.  It is
	// commonly used by selinux
	ProcessLabel string `json:"process_label,omitempty"`

	// Rlimits specifies the resource limits, such as max open files, to set in the container
	// If Rlimits are not set, the container will inherit rlimits from the parent process
	Rlimits []Rlimit `json:"rlimits,omitempty"`

	// OomScoreAdj specifies the adjustment to be made by the kernel when calculating oom scores
	// for a process. Valid values are between the range [-1000, '1000'], where processes with
	// higher scores are preferred for being killed.
	// More information about kernel oom score calculation here: https://lwn.net/Articles/317814/
	OomScoreAdj int `json:"oom_score_adj"`

	// AdditionalGroups specifies the gids that should be added to supplementary groups
	// in addition to those that the user belongs to.
	AdditionalGroups []string `json:"additional_groups"`

	// UidMappings is an array of User ID mappings for User Namespaces
	UidMappings []IDMap `json:"uid_mappings"`

	// GidMappings is an array of Group ID mappings for User Namespaces
	GidMappings []IDMap `json:"gid_mappings"`

	// MaskPaths specifies paths within the container's rootfs to mask over with a bind
	// mount pointing to /dev/null as to prevent reads of the file.
	MaskPaths []string `json:"mask_paths"`

	// ReadonlyPaths specifies paths within the container's rootfs to remount as read-only
	// so that these files prevent any writes.
	ReadonlyPaths []string `json:"readonly_paths"`

	// Sysctl is a map of properties and their values. It is the equivalent of using
	// sysctl -w my.property.name value in Linux.
	Sysctl map[string]string `json:"sysctl"`

	// Seccomp allows actions to be taken whenever a syscall is made within the container.
	// A number of rules are given, each having an action to be taken if a syscall matches it.
	// A default action to be taken if no rules match is also given.
	Seccomp *Seccomp `json:"seccomp"`

	// NoNewPrivileges controls whether processes in the container can gain additional privileges.
	NoNewPrivileges bool `json:"no_new_privileges,omitempty"`

	// Hooks are a collection of actions to perform at various container lifecycle events.
	// CommandHooks are serialized to JSON, but other hooks are not.
	Hooks *Hooks

	// Version is the version of opencontainer specification that is supported.
	Version string `json:"version"`

	// Labels are user defined metadata that is stored in the config and populated on the state
	Labels []string `json:"labels"`
}

type Hooks struct {
	// Prestart commands are executed after the container namespaces are created,
	// but before the user supplied command is executed from init.
	Prestart []Hook

	// Poststart commands are executed after the container init process starts.
	Poststart []Hook

	// Poststop commands are executed after the container init process exits.
	Poststop []Hook
}

func (hooks *Hooks) UnmarshalJSON(b []byte) error {
	var state struct {
		Prestart  []CommandHook
		Poststart []CommandHook
		Poststop  []CommandHook
	}

	if err := json.Unmarshal(b, &state); err != nil {
		return err
	}

	deserialize := func(shooks []CommandHook) (hooks []Hook) {
		for _, shook := range shooks {
			hooks = append(hooks, shook)
		}

		return hooks
	}

	hooks.Prestart = deserialize(state.Prestart)
	hooks.Poststart = deserialize(state.Poststart)
	hooks.Poststop = deserialize(state.Poststop)
	return nil
}

func (hooks Hooks) MarshalJSON() ([]byte, error) {
	serialize := func(hooks []Hook) (serializableHooks []CommandHook) {
		for _, hook := range hooks {
			switch chook := hook.(type) {
			case CommandHook:
				serializableHooks = append(serializableHooks, chook)
			default:
				logrus.Warnf("cannot serialize hook of type %T, skipping", hook)
			}
		}

		return serializableHooks
	}

	return json.Marshal(map[string]interface{}{
		"prestart":  serialize(hooks.Prestart),
		"poststart": serialize(hooks.Poststart),
		"poststop":  serialize(hooks.Poststop),
	})
}

// HookState is the payload provided to a hook on execution.
type HookState struct {
	Version string `json:"version"`
	ID      string `json:"id"`
	Pid     int    `json:"pid"`
	Root    string `json:"root"`
}

type Hook interface {
	// Run executes the hook with the provided state.
	Run(HookState) error
}

// NewFunctionHooks will call the provided function when the hook is run.
func NewFunctionHook(f func(HookState) error) FuncHook {
	return FuncHook{
		run: f,
	}
}

type FuncHook struct {
	run func(HookState) error
}

func (f FuncHook) Run(s HookState) error {
	return f.run(s)
}

type Command struct {
	Path    string         `json:"path"`
	Args    []string       `json:"args"`
	Env     []string       `json:"env"`
	Dir     string         `json:"dir"`
	Timeout *time.Duration `json:"timeout"`
}

// NewCommandHooks will execute the provided command when the hook is run.
func NewCommandHook(cmd Command) CommandHook {
	return CommandHook{
		Command: cmd,
	}
}

type CommandHook struct {
	Command
}

func (c Command) Run(s HookState) error {
	b, err := json.Marshal(s)
	if err != nil {
		return err
	}
	cmd := exec.Cmd{
		Path:  c.Path,
		Args:  c.Args,
		Env:   c.Env,
		Stdin: bytes.NewReader(b),
	}
	errC := make(chan error, 1)
	go func() {
		out, err := cmd.CombinedOutput()
		if err != nil {
			err = fmt.Errorf("%s: %s", err, out)
		}
		errC <- err
	}()
	if c.Timeout != nil {
		select {
		case err := <-errC:
			return err
		case <-time.After(*c.Timeout):
			cmd.Process.Kill()
			cmd.Wait()
			return fmt.Errorf("hook ran past specified timeout of %.1fs", c.Timeout.Seconds())
		}
	}
	return <-errC
}
<|endoftext|>