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

//go:build windows
// +build windows

package windowsperfcountersreceiver

import (
	"context"
	"path/filepath"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.opentelemetry.io/collector/component/componenttest"
	"go.opentelemetry.io/collector/pdata/pcommon"
	"go.opentelemetry.io/collector/pdata/pmetric"
	"go.opentelemetry.io/collector/receiver/scraperhelper"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"go.uber.org/zap/zaptest/observer"

	"github.com/open-telemetry/opentelemetry-collector-contrib/internal/scrapertest"
	"github.com/open-telemetry/opentelemetry-collector-contrib/internal/scrapertest/golden"
	"github.com/open-telemetry/opentelemetry-collector-contrib/pkg/winperfcounters"
)

type mockPerfCounter struct {
	counterValues []winperfcounters.CounterValue
	metricRep     MetricRep
	path          string
	scrapeErr     error
	closeErr      error
}

func (w *mockPerfCounter) Path() string {
	return w.path
}

func (w *mockPerfCounter) ScrapeData() ([]winperfcounters.CounterValue, error) {
	return w.counterValues, w.scrapeErr
}

func (w *mockPerfCounter) Close() error {
	return w.closeErr
}

func mockPerfCounterFactory(mpc mockPerfCounter) newWatcherFunc {
	return func(string, string, string) (winperfcounters.PerfCounterWatcher, error) {
		return &mpc, nil
	}
}

func Test_WindowsPerfCounterScraper(t *testing.T) {
	type testCase struct {
		name string
		cfg  *Config

		mockCounterPath string
		startMessage    string
		startErr        string

		expectedMetricPath string
	}

	defaultConfig := createDefaultConfig().(*Config)

	testCases := []testCase{
		{
			name: "Standard",
			cfg: &Config{
				MetricMetaData: map[string]MetricConfig{
					"cpu.idle": {
						Description: "percentage of time CPU is idle.",
						Unit:        "%",
						Gauge:       GaugeMetric{},
					},
					"bytes.committed": {
						Description: "number of bytes committed to memory",
						Unit:        "By",
						Gauge:       GaugeMetric{},
					},
					"processor.time": {
						Description: "amount of time processor is busy",
						Unit:        "%",
						Gauge:       GaugeMetric{},
					},
				},
				PerfCounters: []ObjectConfig{
					{Object: "Memory", Counters: []CounterConfig{{Name: "Committed Bytes", MetricRep: MetricRep{Name: "bytes.committed"}}}},
					{Object: "Processor", Instances: []string{"*"}, Counters: []CounterConfig{{Name: "% Idle Time", MetricRep: MetricRep{Name: "cpu.idle"}}}},
					{Object: "Processor", Instances: []string{"1", "2"}, Counters: []CounterConfig{{Name: "% Processor Time", MetricRep: MetricRep{Name: "processor.time"}}}},
				},
				ScraperControllerSettings: scraperhelper.ScraperControllerSettings{CollectionInterval: time.Minute},
			},
			expectedMetricPath: filepath.Join("testdata", "scraper", "standard.json"),
		},
		{
			name: "SumMetric",
			cfg: &Config{
				MetricMetaData: map[string]MetricConfig{
					"bytes.committed": {
						Description: "number of bytes committed to memory",
						Unit:        "By",
						Sum:         SumMetric{},
					},
				},
				PerfCounters: []ObjectConfig{
					{Object: "Memory", Counters: []CounterConfig{{Name: "Committed Bytes", MetricRep: MetricRep{Name: "bytes.committed"}}}},
				},
				ScraperControllerSettings: scraperhelper.ScraperControllerSettings{CollectionInterval: time.Minute},
			},
			expectedMetricPath: filepath.Join("testdata", "scraper", "sum_metric.json"),
		},
		{
			name: "NoMetricDefinition",
			cfg: &Config{
				PerfCounters: []ObjectConfig{
					{Object: "Memory", Counters: []CounterConfig{{Name: "Committed Bytes"}}},
				},
				ScraperControllerSettings: scraperhelper.ScraperControllerSettings{CollectionInterval: time.Minute},
			},
			expectedMetricPath: filepath.Join("testdata", "scraper", "no_metric_def.json"),
		},
		{
			name: "InvalidCounter",
			cfg: &Config{
				PerfCounters: []ObjectConfig{
					{
						Object:   "Memory",
						Counters: []CounterConfig{{Name: "Committed Bytes", MetricRep: MetricRep{Name: "Committed Bytes"}}},
					},
					{
						Object:   "Invalid Object",
						Counters: []CounterConfig{{Name: "Invalid Counter", MetricRep: MetricRep{Name: "invalid"}}},
					},
				},
				ScraperControllerSettings: scraperhelper.ScraperControllerSettings{CollectionInterval: time.Minute},
			},
			startMessage: "some performance counters could not be initialized",
			startErr:     "failed to create perf counter with path \\Invalid Object\\Invalid Counter: The specified object was not found on the computer.\r\n",
		},
	}

	for _, test := range testCases {
		t.Run(test.name, func(t *testing.T) {
			cfg := test.cfg
			if cfg == nil {
				cfg = defaultConfig
			}

			core, obs := observer.New(zapcore.WarnLevel)
			logger := zap.New(core)
			settings := componenttest.NewNopTelemetrySettings()
			settings.Logger = logger
			scraper := newScraper(cfg, settings)

			err := scraper.start(context.Background(), componenttest.NewNopHost())
			if test.startErr != "" {
				require.Equal(t, 1, obs.Len())
				log := obs.All()[0]
				assert.Equal(t, log.Level, zapcore.WarnLevel)
				assert.Equal(t, test.startMessage, log.Message)
				assert.Equal(t, "error", log.Context[0].Key)
				assert.EqualError(t, log.Context[0].Interface.(error), test.startErr)
				return
			}
			require.NoError(t, err)

			actualMetrics, err := scraper.scrape(context.Background())
			require.NoError(t, err)

			err = scraper.shutdown(context.Background())

			require.NoError(t, err)
			expectedMetrics, err := golden.ReadMetrics(test.expectedMetricPath)
			scrapertest.CompareMetrics(expectedMetrics, actualMetrics, scrapertest.IgnoreMetricValues())
			require.NoError(t, err)
		})
	}
}

func TestInitWatchers(t *testing.T) {
	testCases := []struct {
		name          string
		cfgs          []ObjectConfig
		expectedErr   string
		expectedPaths []string
	}{
		{
			name: "basicPath",
			cfgs: []ObjectConfig{
				{
					Object:   "Memory",
					Counters: []CounterConfig{{Name: "Committed Bytes"}},
				},
			},
			expectedPaths: []string{"\\Memory\\Committed Bytes"},
		},
		{
			name: "multiplePaths",
			cfgs: []ObjectConfig{
				{
					Object:   "Memory",
					Counters: []CounterConfig{{Name: "Committed Bytes"}},
				},
				{
					Object:   "Memory",
					Counters: []CounterConfig{{Name: "Available Bytes"}},
				},
			},
			expectedPaths: []string{"\\Memory\\Committed Bytes", "\\Memory\\Available Bytes"},
		},
		{
			name: "multipleIndividualCounters",
			cfgs: []ObjectConfig{
				{
					Object: "Memory",
					Counters: []CounterConfig{
						{Name: "Committed Bytes"},
						{Name: "Available Bytes"},
					},
				},
				{
					Object:   "Memory",
					Counters: []CounterConfig{},
				},
			},
			expectedPaths: []string{"\\Memory\\Committed Bytes", "\\Memory\\Available Bytes"},
		},
		{
			name: "invalidCounter",
			cfgs: []ObjectConfig{
				{
					Object:   "Broken",
					Counters: []CounterConfig{{Name: "Broken Counter"}},
				},
			},

			expectedErr: "failed to create perf counter with path \\Broken\\Broken Counter: The specified object was not found on the computer.\r\n",
		},
		{
			name: "multipleInvalidCounters",
			cfgs: []ObjectConfig{
				{
					Object:   "Broken",
					Counters: []CounterConfig{{Name: "Broken Counter"}},
				},
				{
					Object:   "Broken part 2",
					Counters: []CounterConfig{{Name: "Broken again"}},
				},
			},
			expectedErr: "failed to create perf counter with path \\Broken\\Broken Counter: The specified object was not found on the computer.\r\n; failed to create perf counter with path \\Broken part 2\\Broken again: The specified object was not found on the computer.\r\n",
		},
	}

	for _, test := range testCases {
		t.Run(test.name, func(t *testing.T) {
			s := &scraper{cfg: &Config{PerfCounters: test.cfgs}, newWatcher: winperfcounters.NewWatcher}
			watchers, errs := s.initWatchers()
			if test.expectedErr != "" {
				require.EqualError(t, errs, test.expectedErr)
			} else {
				require.NoError(t, errs)
			}
			for i, watcher := range watchers {
				require.Equal(t, test.expectedPaths[i], watcher.Path())
			}
		})
	}
}

func TestScrape(t *testing.T) {
	testCases := []struct {
		name              string
		cfg               Config
		mockCounterValues []winperfcounters.CounterValue
	}{
		{
			name: "metricsWithoutInstance",
			cfg: Config{
				PerfCounters: []ObjectConfig{
					{
						Counters: []CounterConfig{
							{
								MetricRep: MetricRep{
									Name: "metric1",
								},
							},
							{
								MetricRep: MetricRep{
									Name: "metric2",
									Attributes: map[string]string{
										"test.attribute": "test-value",
									},
								},
							},
						},
					},
				},
				MetricMetaData: map[string]MetricConfig{
					"metric1": {Description: "metric1 description", Unit: "1"},
					"metric2": {Description: "metric2 description", Unit: "2"},
				},
			},
			mockCounterValues: []winperfcounters.CounterValue{{Value: 1.0}},
		},
		{
			name: "metricsWithInstance",
			cfg: Config{
				PerfCounters: []ObjectConfig{
					{
						Counters: []CounterConfig{
							{
								MetricRep: MetricRep{
									Name: "metric1",
								},
							},
							{
								MetricRep: MetricRep{
									Name: "metric2",
									Attributes: map[string]string{
										"test.attribute": "test-value",
									},
								},
							},
						},
					},
				},
				MetricMetaData: map[string]MetricConfig{
					"metric1": {Description: "metric1 description", Unit: "1"},
					"metric2": {Description: "metric2 description", Unit: "2"},
				},
			},
			mockCounterValues: []winperfcounters.CounterValue{{InstanceName: "Test Instance", Value: 1.0}},
		},
	}
	for _, test := range testCases {
		t.Run(test.name, func(t *testing.T) {
			mpc := mockPerfCounter{counterValues: test.mockCounterValues}
			s := &scraper{cfg: &test.cfg, newWatcher: mockPerfCounterFactory(mpc)}
			errs := s.start(context.Background(), componenttest.NewNopHost())
			require.NoError(t, errs)

			m, err := s.scrape(context.Background())
			require.NoError(t, err)
			require.Equal(t, 1, m.ResourceMetrics().Len())
			require.Equal(t, 1, m.ResourceMetrics().At(0).ScopeMetrics().Len())

			metrics := m.ResourceMetrics().At(0).ScopeMetrics().At(0).Metrics()
			metrics.Sort(func(a, b pmetric.Metric) bool {
				return a.Name() < b.Name()
			})
			curMetricsNum := 0
			for _, pc := range test.cfg.PerfCounters {
				for _, counterCfg := range pc.Counters {
					metric := metrics.At(curMetricsNum)
					assert.Equal(t, counterCfg.MetricRep.Name, metric.Name())
					metricData := test.cfg.MetricMetaData[counterCfg.MetricRep.Name]
					assert.Equal(t, metricData.Description, metric.Description())
					assert.Equal(t, metricData.Unit, metric.Unit())
					dps := metric.Gauge().DataPoints()
					assert.Equal(t, len(test.mockCounterValues), dps.Len())
					for dpIdx, val := range test.mockCounterValues {
						assert.Equal(t, val.Value, dps.At(dpIdx).DoubleValue())
						expectedAttributeLen := len(counterCfg.MetricRep.Attributes)
						if val.InstanceName != "" {
							expectedAttributeLen++
						}
						assert.Equal(t, expectedAttributeLen, dps.At(dpIdx).Attributes().Len())
						dps.At(dpIdx).Attributes().Range(func(k string, v pcommon.Value) bool {
							if k == instanceLabelName {
								assert.Equal(t, val.InstanceName, v.Str())
								return true
							}
							assert.Equal(t, counterCfg.MetricRep.Attributes[k], v.Str())
							return true
						})
					}
					curMetricsNum++
				}
			}
		})
	}
}
<commit_msg>[receiver/winperfcounters] Skip Flaky Test_WindowsPerfCounterScraper (#12956)<commit_after>// Copyright The OpenTelemetry Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//       http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//go:build windows
// +build windows

package windowsperfcountersreceiver

import (
	"context"
	"path/filepath"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.opentelemetry.io/collector/component/componenttest"
	"go.opentelemetry.io/collector/pdata/pcommon"
	"go.opentelemetry.io/collector/pdata/pmetric"
	"go.opentelemetry.io/collector/receiver/scraperhelper"
	"go.uber.org/zap"
	"go.uber.org/zap/zapcore"
	"go.uber.org/zap/zaptest/observer"

	"github.com/open-telemetry/opentelemetry-collector-contrib/internal/scrapertest"
	"github.com/open-telemetry/opentelemetry-collector-contrib/internal/scrapertest/golden"
	"github.com/open-telemetry/opentelemetry-collector-contrib/pkg/winperfcounters"
)

type mockPerfCounter struct {
	counterValues []winperfcounters.CounterValue
	metricRep     MetricRep
	path          string
	scrapeErr     error
	closeErr      error
}

func (w *mockPerfCounter) Path() string {
	return w.path
}

func (w *mockPerfCounter) ScrapeData() ([]winperfcounters.CounterValue, error) {
	return w.counterValues, w.scrapeErr
}

func (w *mockPerfCounter) Close() error {
	return w.closeErr
}

func mockPerfCounterFactory(mpc mockPerfCounter) newWatcherFunc {
	return func(string, string, string) (winperfcounters.PerfCounterWatcher, error) {
		return &mpc, nil
	}
}

func Test_WindowsPerfCounterScraper(t *testing.T) {
	t.Skip("Flaky Test - See https://github.com/open-telemetry/opentelemetry-collector-contrib/issues/11445")
	type testCase struct {
		name string
		cfg  *Config

		mockCounterPath string
		startMessage    string
		startErr        string

		expectedMetricPath string
	}

	defaultConfig := createDefaultConfig().(*Config)

	testCases := []testCase{
		{
			name: "Standard",
			cfg: &Config{
				MetricMetaData: map[string]MetricConfig{
					"cpu.idle": {
						Description: "percentage of time CPU is idle.",
						Unit:        "%",
						Gauge:       GaugeMetric{},
					},
					"bytes.committed": {
						Description: "number of bytes committed to memory",
						Unit:        "By",
						Gauge:       GaugeMetric{},
					},
					"processor.time": {
						Description: "amount of time processor is busy",
						Unit:        "%",
						Gauge:       GaugeMetric{},
					},
				},
				PerfCounters: []ObjectConfig{
					{Object: "Memory", Counters: []CounterConfig{{Name: "Committed Bytes", MetricRep: MetricRep{Name: "bytes.committed"}}}},
					{Object: "Processor", Instances: []string{"*"}, Counters: []CounterConfig{{Name: "% Idle Time", MetricRep: MetricRep{Name: "cpu.idle"}}}},
					{Object: "Processor", Instances: []string{"1", "2"}, Counters: []CounterConfig{{Name: "% Processor Time", MetricRep: MetricRep{Name: "processor.time"}}}},
				},
				ScraperControllerSettings: scraperhelper.ScraperControllerSettings{CollectionInterval: time.Minute},
			},
			expectedMetricPath: filepath.Join("testdata", "scraper", "standard.json"),
		},
		{
			name: "SumMetric",
			cfg: &Config{
				MetricMetaData: map[string]MetricConfig{
					"bytes.committed": {
						Description: "number of bytes committed to memory",
						Unit:        "By",
						Sum:         SumMetric{},
					},
				},
				PerfCounters: []ObjectConfig{
					{Object: "Memory", Counters: []CounterConfig{{Name: "Committed Bytes", MetricRep: MetricRep{Name: "bytes.committed"}}}},
				},
				ScraperControllerSettings: scraperhelper.ScraperControllerSettings{CollectionInterval: time.Minute},
			},
			expectedMetricPath: filepath.Join("testdata", "scraper", "sum_metric.json"),
		},
		{
			name: "NoMetricDefinition",
			cfg: &Config{
				PerfCounters: []ObjectConfig{
					{Object: "Memory", Counters: []CounterConfig{{Name: "Committed Bytes"}}},
				},
				ScraperControllerSettings: scraperhelper.ScraperControllerSettings{CollectionInterval: time.Minute},
			},
			expectedMetricPath: filepath.Join("testdata", "scraper", "no_metric_def.json"),
		},
		{
			name: "InvalidCounter",
			cfg: &Config{
				PerfCounters: []ObjectConfig{
					{
						Object:   "Memory",
						Counters: []CounterConfig{{Name: "Committed Bytes", MetricRep: MetricRep{Name: "Committed Bytes"}}},
					},
					{
						Object:   "Invalid Object",
						Counters: []CounterConfig{{Name: "Invalid Counter", MetricRep: MetricRep{Name: "invalid"}}},
					},
				},
				ScraperControllerSettings: scraperhelper.ScraperControllerSettings{CollectionInterval: time.Minute},
			},
			startMessage: "some performance counters could not be initialized",
			startErr:     "failed to create perf counter with path \\Invalid Object\\Invalid Counter: The specified object was not found on the computer.\r\n",
		},
	}

	for _, test := range testCases {
		t.Run(test.name, func(t *testing.T) {
			cfg := test.cfg
			if cfg == nil {
				cfg = defaultConfig
			}

			core, obs := observer.New(zapcore.WarnLevel)
			logger := zap.New(core)
			settings := componenttest.NewNopTelemetrySettings()
			settings.Logger = logger
			scraper := newScraper(cfg, settings)

			err := scraper.start(context.Background(), componenttest.NewNopHost())
			if test.startErr != "" {
				require.Equal(t, 1, obs.Len())
				log := obs.All()[0]
				assert.Equal(t, log.Level, zapcore.WarnLevel)
				assert.Equal(t, test.startMessage, log.Message)
				assert.Equal(t, "error", log.Context[0].Key)
				assert.EqualError(t, log.Context[0].Interface.(error), test.startErr)
				return
			}
			require.NoError(t, err)

			actualMetrics, err := scraper.scrape(context.Background())
			require.NoError(t, err)

			err = scraper.shutdown(context.Background())

			require.NoError(t, err)
			expectedMetrics, err := golden.ReadMetrics(test.expectedMetricPath)
			scrapertest.CompareMetrics(expectedMetrics, actualMetrics, scrapertest.IgnoreMetricValues())
			require.NoError(t, err)
		})
	}
}

func TestInitWatchers(t *testing.T) {
	testCases := []struct {
		name          string
		cfgs          []ObjectConfig
		expectedErr   string
		expectedPaths []string
	}{
		{
			name: "basicPath",
			cfgs: []ObjectConfig{
				{
					Object:   "Memory",
					Counters: []CounterConfig{{Name: "Committed Bytes"}},
				},
			},
			expectedPaths: []string{"\\Memory\\Committed Bytes"},
		},
		{
			name: "multiplePaths",
			cfgs: []ObjectConfig{
				{
					Object:   "Memory",
					Counters: []CounterConfig{{Name: "Committed Bytes"}},
				},
				{
					Object:   "Memory",
					Counters: []CounterConfig{{Name: "Available Bytes"}},
				},
			},
			expectedPaths: []string{"\\Memory\\Committed Bytes", "\\Memory\\Available Bytes"},
		},
		{
			name: "multipleIndividualCounters",
			cfgs: []ObjectConfig{
				{
					Object: "Memory",
					Counters: []CounterConfig{
						{Name: "Committed Bytes"},
						{Name: "Available Bytes"},
					},
				},
				{
					Object:   "Memory",
					Counters: []CounterConfig{},
				},
			},
			expectedPaths: []string{"\\Memory\\Committed Bytes", "\\Memory\\Available Bytes"},
		},
		{
			name: "invalidCounter",
			cfgs: []ObjectConfig{
				{
					Object:   "Broken",
					Counters: []CounterConfig{{Name: "Broken Counter"}},
				},
			},

			expectedErr: "failed to create perf counter with path \\Broken\\Broken Counter: The specified object was not found on the computer.\r\n",
		},
		{
			name: "multipleInvalidCounters",
			cfgs: []ObjectConfig{
				{
					Object:   "Broken",
					Counters: []CounterConfig{{Name: "Broken Counter"}},
				},
				{
					Object:   "Broken part 2",
					Counters: []CounterConfig{{Name: "Broken again"}},
				},
			},
			expectedErr: "failed to create perf counter with path \\Broken\\Broken Counter: The specified object was not found on the computer.\r\n; failed to create perf counter with path \\Broken part 2\\Broken again: The specified object was not found on the computer.\r\n",
		},
	}

	for _, test := range testCases {
		t.Run(test.name, func(t *testing.T) {
			s := &scraper{cfg: &Config{PerfCounters: test.cfgs}, newWatcher: winperfcounters.NewWatcher}
			watchers, errs := s.initWatchers()
			if test.expectedErr != "" {
				require.EqualError(t, errs, test.expectedErr)
			} else {
				require.NoError(t, errs)
			}
			for i, watcher := range watchers {
				require.Equal(t, test.expectedPaths[i], watcher.Path())
			}
		})
	}
}

func TestScrape(t *testing.T) {
	testCases := []struct {
		name              string
		cfg               Config
		mockCounterValues []winperfcounters.CounterValue
	}{
		{
			name: "metricsWithoutInstance",
			cfg: Config{
				PerfCounters: []ObjectConfig{
					{
						Counters: []CounterConfig{
							{
								MetricRep: MetricRep{
									Name: "metric1",
								},
							},
							{
								MetricRep: MetricRep{
									Name: "metric2",
									Attributes: map[string]string{
										"test.attribute": "test-value",
									},
								},
							},
						},
					},
				},
				MetricMetaData: map[string]MetricConfig{
					"metric1": {Description: "metric1 description", Unit: "1"},
					"metric2": {Description: "metric2 description", Unit: "2"},
				},
			},
			mockCounterValues: []winperfcounters.CounterValue{{Value: 1.0}},
		},
		{
			name: "metricsWithInstance",
			cfg: Config{
				PerfCounters: []ObjectConfig{
					{
						Counters: []CounterConfig{
							{
								MetricRep: MetricRep{
									Name: "metric1",
								},
							},
							{
								MetricRep: MetricRep{
									Name: "metric2",
									Attributes: map[string]string{
										"test.attribute": "test-value",
									},
								},
							},
						},
					},
				},
				MetricMetaData: map[string]MetricConfig{
					"metric1": {Description: "metric1 description", Unit: "1"},
					"metric2": {Description: "metric2 description", Unit: "2"},
				},
			},
			mockCounterValues: []winperfcounters.CounterValue{{InstanceName: "Test Instance", Value: 1.0}},
		},
	}
	for _, test := range testCases {
		t.Run(test.name, func(t *testing.T) {
			mpc := mockPerfCounter{counterValues: test.mockCounterValues}
			s := &scraper{cfg: &test.cfg, newWatcher: mockPerfCounterFactory(mpc)}
			errs := s.start(context.Background(), componenttest.NewNopHost())
			require.NoError(t, errs)

			m, err := s.scrape(context.Background())
			require.NoError(t, err)
			require.Equal(t, 1, m.ResourceMetrics().Len())
			require.Equal(t, 1, m.ResourceMetrics().At(0).ScopeMetrics().Len())

			metrics := m.ResourceMetrics().At(0).ScopeMetrics().At(0).Metrics()
			metrics.Sort(func(a, b pmetric.Metric) bool {
				return a.Name() < b.Name()
			})
			curMetricsNum := 0
			for _, pc := range test.cfg.PerfCounters {
				for _, counterCfg := range pc.Counters {
					metric := metrics.At(curMetricsNum)
					assert.Equal(t, counterCfg.MetricRep.Name, metric.Name())
					metricData := test.cfg.MetricMetaData[counterCfg.MetricRep.Name]
					assert.Equal(t, metricData.Description, metric.Description())
					assert.Equal(t, metricData.Unit, metric.Unit())
					dps := metric.Gauge().DataPoints()
					assert.Equal(t, len(test.mockCounterValues), dps.Len())
					for dpIdx, val := range test.mockCounterValues {
						assert.Equal(t, val.Value, dps.At(dpIdx).DoubleValue())
						expectedAttributeLen := len(counterCfg.MetricRep.Attributes)
						if val.InstanceName != "" {
							expectedAttributeLen++
						}
						assert.Equal(t, expectedAttributeLen, dps.At(dpIdx).Attributes().Len())
						dps.At(dpIdx).Attributes().Range(func(k string, v pcommon.Value) bool {
							if k == instanceLabelName {
								assert.Equal(t, val.InstanceName, v.Str())
								return true
							}
							assert.Equal(t, counterCfg.MetricRep.Attributes[k], v.Str())
							return true
						})
					}
					curMetricsNum++
				}
			}
		})
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2022 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 zanzibar

import (
	"context"
	"time"

	jsonschema "github.com/mcuadros/go-jsonschema-generator"
	"github.com/uber-go/tally"
)

const (
	middlewareRequestLatencyTag  = "middleware.request.latency"
	middlewareResponseLatencyTag = "middleware.response.latency"
	middlewareRequestStatusTag   = "middleware.request.status"
)

// MiddlewareStack is a stack of Middleware Handlers that can be invoked as an Handle.
// MiddlewareStack middlewares are evaluated for requests in the order that they are added to the stack
// followed by the underlying HandlerFn. The middleware responses are then executed in reverse.
type MiddlewareStack struct {
	middlewares []MiddlewareHandle
	handle      HandlerFn
}

// NewStack returns a new MiddlewareStack instance with no middleware preconfigured.
func NewStack(middlewares []MiddlewareHandle,
	handle HandlerFn) *MiddlewareStack {
	return &MiddlewareStack{
		handle:      handle,
		middlewares: middlewares,
	}
}

// Middlewares returns a list of all the handlers in the current MiddlewareStack.
func (m *MiddlewareStack) Middlewares() []MiddlewareHandle {
	return m.middlewares
}

// MiddlewareHandle used to define middleware
type MiddlewareHandle interface {
	// implement HandleRequest for your middleware. Return false
	// if the handler writes to the response body.
	HandleRequest(
		ctx context.Context,
		req *ServerHTTPRequest,
		res *ServerHTTPResponse,
		shared SharedState,
	) (context.Context, bool)
	// implement HandleResponse for your middleware. Return false
	// if the handler writes to the response body.
	HandleResponse(
		ctx context.Context,
		res *ServerHTTPResponse,
		shared SharedState,
	) context.Context
	// return any shared state for this middleware.
	JSONSchema() *jsonschema.Document
	Name() string
}

// SharedState used to access other middlewares in the chain.
type SharedState struct {
	middlewareDict map[string]interface{}
}

// NewSharedState constructs a ShardState
func NewSharedState(middlewares []MiddlewareHandle) SharedState {
	sharedState := SharedState{}
	sharedState.middlewareDict = make(map[string]interface{})

	for i := 0; i < len(middlewares); i++ {
		sharedState.middlewareDict[middlewares[i].Name()] = nil
	}
	return sharedState
}

// GetState returns the state from a different middleware
func (s SharedState) GetState(name string) interface{} {
	return s.middlewareDict[name]
}

// SetState sets value of a middleware shared state
func (s SharedState) SetState(m MiddlewareHandle, state interface{}) {
	s.middlewareDict[m.Name()] = state
}

// Handle executes the middlewares in a stack and underlying handler.
func (m *MiddlewareStack) Handle(
	ctx context.Context,
	req *ServerHTTPRequest,
	res *ServerHTTPResponse) context.Context {

	shared := NewSharedState(m.middlewares)
	middlewareRequestStartTime := time.Now()

	for i := 0; i < len(m.middlewares); i++ {
		ctx, ok := m.middlewares[i].HandleRequest(ctx, req, res, shared)
		// If a middleware errors and writes to the response header
		// then abort the rest of the stack and evaluate the response
		// handlers for the middlewares seen so far.
		if ok == false {
			//record latency for middlewares requests in unsuccesful case as the middleware requests calls are terminated
			m.recordLatency(middlewareRequestLatencyTag, middlewareRequestStartTime, req.scope)

			middlewareResponseStartTime := time.Now() // start the timer for middleware responses
			for j := i; j >= 0; j-- {
				m.middlewares[j].HandleResponse(ctx, res, shared)
			}
			//record latency for middlewares responses in unsuccesful case
			m.recordLatency(middlewareResponseLatencyTag, middlewareResponseStartTime, req.scope)

			//for error metrics only emit when there is gateway error and not request error
			// the percentage can be calculated via error_count/total_request
			if res.pendingStatusCode >= 500 {
				m.emitAvailabilityError(middlewareRequestStatusTag, req.scope)
			}
			return ctx
		}
	}
	// record latency for middlewares requests in successful case
	m.recordLatency(middlewareRequestLatencyTag, middlewareRequestStartTime, req.scope)

	ctx = m.handle(ctx, req, res)

	middlewareResponseStartTime := time.Now()
	for i := len(m.middlewares) - 1; i >= 0; i-- {
		m.middlewares[i].HandleResponse(ctx, res, shared)
	}
	// record latency for middlewares responses in successful case
	m.recordLatency(middlewareResponseLatencyTag, middlewareResponseStartTime, req.scope)
	return ctx
}

// recordLatency measures the latency as per the tagName and start time given.
func (m *MiddlewareStack) recordLatency(tagName string, startTime time.Time, scope tally.Scope) {
	elapsed := time.Now().Sub(startTime)
	scope.Timer(tagName).Record(elapsed)
	scope.Histogram(tagName, tally.DefaultBuckets).RecordDuration(elapsed)
}

// emitAvailability is used to increment the error counter for a particular tagName.
func (m *MiddlewareStack) emitAvailabilityError(tagName string, scope tally.Scope) {
	tagged := scope.Tagged(map[string]string{
		scopeTagStatus: "error",
	})
	tagged.Counter(tagName).Inc(1)
}
<commit_msg>fixed comments<commit_after>// Copyright (c) 2022 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 zanzibar

import (
	"context"
	"time"

	jsonschema "github.com/mcuadros/go-jsonschema-generator"
	"github.com/uber-go/tally"
)

const (
	middlewareRequestLatencyTag  = "middleware.request.latency"
	middlewareResponseLatencyTag = "middleware.response.latency"
	middlewareRequestStatusTag   = "middleware.request.status"
)

// MiddlewareStack is a stack of Middleware Handlers that can be invoked as an Handle.
// MiddlewareStack middlewares are evaluated for requests in the order that they are added to the stack
// followed by the underlying HandlerFn. The middleware responses are then executed in reverse.
type MiddlewareStack struct {
	middlewares []MiddlewareHandle
	handle      HandlerFn
}

// NewStack returns a new MiddlewareStack instance with no middleware preconfigured.
func NewStack(middlewares []MiddlewareHandle,
	handle HandlerFn) *MiddlewareStack {
	return &MiddlewareStack{
		handle:      handle,
		middlewares: middlewares,
	}
}

// Middlewares returns a list of all the handlers in the current MiddlewareStack.
func (m *MiddlewareStack) Middlewares() []MiddlewareHandle {
	return m.middlewares
}

// MiddlewareHandle used to define middleware
type MiddlewareHandle interface {
	// implement HandleRequest for your middleware. Return false
	// if the handler writes to the response body.
	HandleRequest(
		ctx context.Context,
		req *ServerHTTPRequest,
		res *ServerHTTPResponse,
		shared SharedState,
	) (context.Context, bool)
	// implement HandleResponse for your middleware. Return false
	// if the handler writes to the response body.
	HandleResponse(
		ctx context.Context,
		res *ServerHTTPResponse,
		shared SharedState,
	) context.Context
	// return any shared state for this middleware.
	JSONSchema() *jsonschema.Document
	Name() string
}

// SharedState used to access other middlewares in the chain.
type SharedState struct {
	middlewareDict map[string]interface{}
}

// NewSharedState constructs a ShardState
func NewSharedState(middlewares []MiddlewareHandle) SharedState {
	sharedState := SharedState{}
	sharedState.middlewareDict = make(map[string]interface{})

	for i := 0; i < len(middlewares); i++ {
		sharedState.middlewareDict[middlewares[i].Name()] = nil
	}
	return sharedState
}

// GetState returns the state from a different middleware
func (s SharedState) GetState(name string) interface{} {
	return s.middlewareDict[name]
}

// SetState sets value of a middleware shared state
func (s SharedState) SetState(m MiddlewareHandle, state interface{}) {
	s.middlewareDict[m.Name()] = state
}

// Handle executes the middlewares in a stack and underlying handler.
func (m *MiddlewareStack) Handle(
	ctx context.Context,
	req *ServerHTTPRequest,
	res *ServerHTTPResponse) context.Context {

	shared := NewSharedState(m.middlewares)
	middlewareRequestStartTime := time.Now()

	for i := 0; i < len(m.middlewares); i++ {
		ctx, ok := m.middlewares[i].HandleRequest(ctx, req, res, shared)
		// If a middleware errors and writes to the response header
		// then abort the rest of the stack and evaluate the response
		// handlers for the middlewares seen so far.
		if ok == false {
			//record latency for middlewares requests in unsuccessful case as the middleware requests calls are terminated
			m.recordLatency(middlewareRequestLatencyTag, middlewareRequestStartTime, req.scope)

			middlewareResponseStartTime := time.Now() // start the timer for middleware responses
			for j := i; j >= 0; j-- {
				m.middlewares[j].HandleResponse(ctx, res, shared)
			}
			//record latency for middlewares responses in unsuccessful case
			m.recordLatency(middlewareResponseLatencyTag, middlewareResponseStartTime, req.scope)

			//for error metrics only emit when there is gateway error and not request error
			// the percentage can be calculated via error_count/total_request
			if res.pendingStatusCode >= 500 {
				m.emitAvailabilityError(middlewareRequestStatusTag, req.scope)
			}
			return ctx
		}
	}
	// record latency for middlewares requests in successful case
	m.recordLatency(middlewareRequestLatencyTag, middlewareRequestStartTime, req.scope)

	ctx = m.handle(ctx, req, res)

	middlewareResponseStartTime := time.Now()
	for i := len(m.middlewares) - 1; i >= 0; i-- {
		m.middlewares[i].HandleResponse(ctx, res, shared)
	}
	// record latency for middlewares responses in successful case
	m.recordLatency(middlewareResponseLatencyTag, middlewareResponseStartTime, req.scope)
	return ctx
}

// recordLatency measures the latency as per the tagName and start time given.
func (m *MiddlewareStack) recordLatency(tagName string, startTime time.Time, scope tally.Scope) {
	elapsed := time.Now().Sub(startTime)
	scope.Timer(tagName).Record(elapsed)
	scope.Histogram(tagName, tally.DefaultBuckets).RecordDuration(elapsed)
}

// emitAvailability is used to increment the error counter for a particular tagName.
func (m *MiddlewareStack) emitAvailabilityError(tagName string, scope tally.Scope) {
	tagged := scope.Tagged(map[string]string{
		scopeTagStatus: "error",
	})
	tagged.Counter(tagName).Inc(1)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"github.com/onsi/ginkgo/config"
	"github.com/onsi/ginkgo/ginkgo/aggregator"
	"github.com/onsi/ginkgo/ginkgo/testsuite"
	"github.com/onsi/ginkgo/remote"
	"github.com/onsi/ginkgo/stenographer"
	"io"
	"os"
	"os/exec"
	"time"
)

type testRunner struct {
	numCPU           int
	parallelStream   bool
	runMagicI        bool
	race             bool
	cover            bool
	executedCommands []*exec.Cmd
	reports          []*bytes.Buffer
}

func newTestRunner(numCPU int, parallelStream bool, runMagicI bool, race bool, cover bool) *testRunner {
	return &testRunner{
		numCPU:           numCPU,
		parallelStream:   parallelStream,
		runMagicI:        runMagicI,
		race:             race,
		cover:            cover,
		executedCommands: []*exec.Cmd{},
		reports:          []*bytes.Buffer{},
	}
}

func (t *testRunner) runSuite(suite *testsuite.TestSuite) bool {
	if t.runMagicI {
		err := t.runGoI(suite)
		if err != nil {
			return false
		}
	}

	var success bool
	if suite.IsGinkgo {
		if t.numCPU > 1 {
			if t.parallelStream {
				success = t.runAndStreamParallelGinkgoSuite(suite)
			} else {
				success = t.runParallelGinkgoSuite(suite)
			}
		} else {
			success = t.runSerialGinkgoSuite(suite)
		}
	} else {
		success = t.runGoTestSuite(suite)
	}

	return success
}

func (t *testRunner) runGoI(suite *testsuite.TestSuite) error {
	args := []string{"test", "-i"}
	if t.race {
		args = append(args, "-race")
	}
	args = append(args, suite.Path)
	cmd := exec.Command("go", args...)
	output, err := cmd.CombinedOutput()
	if err != nil {
		fmt.Printf("go test -i %s failed with:\n\n%s", suite.Path, output)
	}

	return err
}

func (t *testRunner) runParallelGinkgoSuite(suite *testsuite.TestSuite) bool {
	completions := make(chan bool)
	for cpu := 0; cpu < t.numCPU; cpu++ {
		config.GinkgoConfig.ParallelNode = cpu + 1
		config.GinkgoConfig.ParallelTotal = t.numCPU

		args := config.BuildFlagArgs("ginkgo", config.GinkgoConfig, config.DefaultReporterConfig)
		args = append(args, t.commonArgs(suite)...)

		buffer := new(bytes.Buffer)
		t.reports = append(t.reports, buffer)

		go t.runCommand(suite.Path, args, nil, buffer, completions)
	}

	passed := true

	for cpu := 0; cpu < t.numCPU; cpu++ {
		passed = <-completions && passed
	}

	for _, report := range t.reports {
		fmt.Print(report.String())
	}
	os.Stdout.Sync()

	return passed
}

func (t *testRunner) runAndStreamParallelGinkgoSuite(suite *testsuite.TestSuite) bool {
	result := make(chan bool, 0)
	stenographer := stenographer.New(!config.DefaultReporterConfig.NoColor)
	aggregator := aggregator.NewAggregator(t.numCPU, result, config.DefaultReporterConfig, stenographer)

	server, err := remote.NewServer()
	if err != nil {
		panic("Failed to start parallel spec server")
	}

	server.RegisterReporters(aggregator)
	server.Start()

	serverAddress := server.Address()

	completions := make(chan bool)

	for cpu := 0; cpu < t.numCPU; cpu++ {
		config.GinkgoConfig.ParallelNode = cpu + 1
		config.GinkgoConfig.ParallelTotal = t.numCPU

		args := config.BuildFlagArgs("ginkgo", config.GinkgoConfig, config.DefaultReporterConfig)
		args = append(args, t.commonArgs(suite)...)

		env := os.Environ()
		env = append(env, fmt.Sprintf("GINKGO_REMOTE_REPORTING_SERVER=%s", serverAddress))

		buffer := new(bytes.Buffer)
		t.reports = append(t.reports, buffer)

		go t.runCommand(suite.Path, args, env, buffer, completions)
	}

	for cpu := 0; cpu < t.numCPU; cpu++ {
		<-completions
	}

	//all test processes are done, at this point
	//we should be able to wait for the aggregator to tell us that it's done

	var passed = false
	select {
	case passed = <-result:
		//the aggregator is done and can tell us whether or not the suite passed
	case <-time.After(time.Second):
		//the aggregator never got back to us!  something must have gone wrong
		fmt.Println("")
		fmt.Println("")
		fmt.Println("   ----------------------------------------------------------- ")
		fmt.Println("  |                                                           |")
		fmt.Println("  |  Ginkgo timed out waiting for all parallel nodes to end!  |")
		fmt.Println("  |  Here is some salvaged output:                            |")
		fmt.Println("  |                                                           |")
		fmt.Println("   ----------------------------------------------------------- ")
		fmt.Println("")
		fmt.Println("")

		os.Stdout.Sync()

		time.Sleep(time.Second)

		for _, report := range t.reports {
			fmt.Print(report.String())
		}

		os.Stdout.Sync()
	}

	server.Stop()

	return passed
}

func (t *testRunner) runSerialGinkgoSuite(suite *testsuite.TestSuite) bool {
	args := config.BuildFlagArgs("ginkgo", config.GinkgoConfig, config.DefaultReporterConfig)
	args = append(args, t.commonArgs(suite)...)
	return t.runCommand(suite.Path, args, nil, os.Stdout, nil)
}

func (t *testRunner) runGoTestSuite(suite *testsuite.TestSuite) bool {
	args := t.commonArgs(suite)
	return t.runCommand(suite.Path, args, nil, os.Stdout, nil)
}

func (t *testRunner) commonArgs(suite *testsuite.TestSuite) []string {
	args := []string{}
	if t.race {
		args = append(args, "--race")
	}
	if t.cover {
		args = append([]string{"--cover", "--coverprofile=" + suite.PackageName + ".coverprofile"})
	}
	return args
}

func (t *testRunner) runCommand(path string, args []string, env []string, stream io.Writer, completions chan bool) bool {
	args = append([]string{"test", "-v", "-timeout=24h", path}, args...)

	cmd := exec.Command("go", args...)
	cmd.Env = env

	t.executedCommands = append(t.executedCommands, cmd)

	doneStreaming := make(chan bool, 2)
	streamPipe := func(pipe io.ReadCloser) {
		io.Copy(stream, pipe)
		doneStreaming <- true
	}

	stdout, _ := cmd.StdoutPipe()
	stderr, _ := cmd.StderrPipe()
	go streamPipe(stdout)
	go streamPipe(stderr)

	err := cmd.Start()
	if err != nil {
		os.Exit(1)
	}

	<-doneStreaming
	<-doneStreaming

	err = cmd.Wait()
	if completions != nil {
		completions <- (err == nil)
	}
	return err == nil
}

func (t *testRunner) abort(sig os.Signal) {
	for _, cmd := range t.executedCommands {
		if cmd.Process != nil {
			cmd.Process.Signal(sig)
			cmd.Wait()
		}
	}
}
<commit_msg>fixed thread-safety issue in test_runner<commit_after>package main

import (
	"bytes"
	"fmt"
	"github.com/onsi/ginkgo/config"
	"github.com/onsi/ginkgo/ginkgo/aggregator"
	"github.com/onsi/ginkgo/ginkgo/testsuite"
	"github.com/onsi/ginkgo/remote"
	"github.com/onsi/ginkgo/stenographer"
	"io"
	"os"
	"os/exec"
	"sync"
	"time"
)

type testRunner struct {
	numCPU           int
	parallelStream   bool
	runMagicI        bool
	race             bool
	cover            bool
	executedCommands []*exec.Cmd
	reports          []*bytes.Buffer

	lock *sync.Mutex
}

func newTestRunner(numCPU int, parallelStream bool, runMagicI bool, race bool, cover bool) *testRunner {
	return &testRunner{
		numCPU:           numCPU,
		parallelStream:   parallelStream,
		runMagicI:        runMagicI,
		race:             race,
		cover:            cover,
		executedCommands: []*exec.Cmd{},
		reports:          []*bytes.Buffer{},
		lock:             &sync.Mutex{},
	}
}

func (t *testRunner) runSuite(suite *testsuite.TestSuite) bool {
	if t.runMagicI {
		err := t.runGoI(suite)
		if err != nil {
			return false
		}
	}

	var success bool
	if suite.IsGinkgo {
		if t.numCPU > 1 {
			if t.parallelStream {
				success = t.runAndStreamParallelGinkgoSuite(suite)
			} else {
				success = t.runParallelGinkgoSuite(suite)
			}
		} else {
			success = t.runSerialGinkgoSuite(suite)
		}
	} else {
		success = t.runGoTestSuite(suite)
	}

	return success
}

func (t *testRunner) runGoI(suite *testsuite.TestSuite) error {
	args := []string{"test", "-i"}
	if t.race {
		args = append(args, "-race")
	}
	args = append(args, suite.Path)
	cmd := exec.Command("go", args...)
	output, err := cmd.CombinedOutput()
	if err != nil {
		fmt.Printf("go test -i %s failed with:\n\n%s", suite.Path, output)
	}

	return err
}

func (t *testRunner) runParallelGinkgoSuite(suite *testsuite.TestSuite) bool {
	completions := make(chan bool)
	for cpu := 0; cpu < t.numCPU; cpu++ {
		config.GinkgoConfig.ParallelNode = cpu + 1
		config.GinkgoConfig.ParallelTotal = t.numCPU

		args := config.BuildFlagArgs("ginkgo", config.GinkgoConfig, config.DefaultReporterConfig)
		args = append(args, t.commonArgs(suite)...)

		buffer := new(bytes.Buffer)
		t.reports = append(t.reports, buffer)

		go t.runCommand(suite.Path, args, nil, buffer, completions)
	}

	passed := true

	for cpu := 0; cpu < t.numCPU; cpu++ {
		passed = <-completions && passed
	}

	for _, report := range t.reports {
		fmt.Print(report.String())
	}
	os.Stdout.Sync()

	return passed
}

func (t *testRunner) runAndStreamParallelGinkgoSuite(suite *testsuite.TestSuite) bool {
	result := make(chan bool, 0)
	stenographer := stenographer.New(!config.DefaultReporterConfig.NoColor)
	aggregator := aggregator.NewAggregator(t.numCPU, result, config.DefaultReporterConfig, stenographer)

	server, err := remote.NewServer()
	if err != nil {
		panic("Failed to start parallel spec server")
	}

	server.RegisterReporters(aggregator)
	server.Start()

	serverAddress := server.Address()

	completions := make(chan bool)

	for cpu := 0; cpu < t.numCPU; cpu++ {
		config.GinkgoConfig.ParallelNode = cpu + 1
		config.GinkgoConfig.ParallelTotal = t.numCPU

		args := config.BuildFlagArgs("ginkgo", config.GinkgoConfig, config.DefaultReporterConfig)
		args = append(args, t.commonArgs(suite)...)

		env := os.Environ()
		env = append(env, fmt.Sprintf("GINKGO_REMOTE_REPORTING_SERVER=%s", serverAddress))

		buffer := new(bytes.Buffer)
		t.reports = append(t.reports, buffer)

		go t.runCommand(suite.Path, args, env, buffer, completions)
	}

	for cpu := 0; cpu < t.numCPU; cpu++ {
		<-completions
	}

	//all test processes are done, at this point
	//we should be able to wait for the aggregator to tell us that it's done

	var passed = false
	select {
	case passed = <-result:
		//the aggregator is done and can tell us whether or not the suite passed
	case <-time.After(time.Second):
		//the aggregator never got back to us!  something must have gone wrong
		fmt.Println("")
		fmt.Println("")
		fmt.Println("   ----------------------------------------------------------- ")
		fmt.Println("  |                                                           |")
		fmt.Println("  |  Ginkgo timed out waiting for all parallel nodes to end!  |")
		fmt.Println("  |  Here is some salvaged output:                            |")
		fmt.Println("  |                                                           |")
		fmt.Println("   ----------------------------------------------------------- ")
		fmt.Println("")
		fmt.Println("")

		os.Stdout.Sync()

		time.Sleep(time.Second)

		for _, report := range t.reports {
			fmt.Print(report.String())
		}

		os.Stdout.Sync()
	}

	server.Stop()

	return passed
}

func (t *testRunner) runSerialGinkgoSuite(suite *testsuite.TestSuite) bool {
	args := config.BuildFlagArgs("ginkgo", config.GinkgoConfig, config.DefaultReporterConfig)
	args = append(args, t.commonArgs(suite)...)
	return t.runCommand(suite.Path, args, nil, os.Stdout, nil)
}

func (t *testRunner) runGoTestSuite(suite *testsuite.TestSuite) bool {
	args := t.commonArgs(suite)
	return t.runCommand(suite.Path, args, nil, os.Stdout, nil)
}

func (t *testRunner) commonArgs(suite *testsuite.TestSuite) []string {
	args := []string{}
	if t.race {
		args = append(args, "--race")
	}
	if t.cover {
		args = append([]string{"--cover", "--coverprofile=" + suite.PackageName + ".coverprofile"})
	}
	return args
}

func (t *testRunner) runCommand(path string, args []string, env []string, stream io.Writer, completions chan bool) bool {

	args = append([]string{"test", "-v", "-timeout=24h", path}, args...)

	cmd := exec.Command("go", args...)
	cmd.Env = env

	t.lock.Lock()
	t.executedCommands = append(t.executedCommands, cmd)
	t.lock.Unlock()

	doneStreaming := make(chan bool, 2)
	streamPipe := func(pipe io.ReadCloser) {
		io.Copy(stream, pipe)
		doneStreaming <- true
	}

	stdout, _ := cmd.StdoutPipe()
	stderr, _ := cmd.StderrPipe()
	go streamPipe(stdout)
	go streamPipe(stderr)

	err := cmd.Start()
	if err != nil {
		os.Exit(1)
	}

	<-doneStreaming
	<-doneStreaming

	err = cmd.Wait()

	t.lock.Lock()
	//delete the command that just finished executing
	for commandIndex, executedCommand := range t.executedCommands {
		if executedCommand == cmd {
			t.executedCommands[commandIndex] = t.executedCommands[len(t.executedCommands)-1]
			t.executedCommands = t.executedCommands[0 : len(t.executedCommands)-1]
			break
		}
	}
	t.lock.Unlock()

	if completions != nil {
		completions <- (err == nil)
	}
	return err == nil
}

func (t *testRunner) abort(sig os.Signal) {
	t.lock.Lock()
	for _, cmd := range t.executedCommands {
		if cmd.Process != nil {
			cmd.Process.Signal(sig)
			cmd.Wait()
		}
	}
	t.lock.Unlock()
}
<|endoftext|>


<commit_before>package generator

import (
	"encoding/xml"
	"fmt"
	"io/ioutil"
	"os"
	"socialapi/config"
	"socialapi/workers/helper"
	"socialapi/workers/sitemap/common"
	"socialapi/workers/sitemap/models"

	"github.com/garyburd/redigo/redis"
	"github.com/koding/logging"
	"github.com/robfig/cron"
)

type Controller struct {
	log          logging.Logger
	fileSelector FileSelector
	fileName     string
}

const (
	// TODO change this later
	SCHEDULE = "0 0-59/30 * * * *"
)

var (
	cronJob *cron.Cron
)

func New(log logging.Logger) (*Controller, error) {
	c := &Controller{
		log:          log,
		fileSelector: CachedFileSelector{},
	}

	return c, nil
}

func (c *Controller) initCron() error {
	cronJob := cron.New()
	if err := cronJob.AddFunc(SCHEDULE, c.generate); err != nil {
		return err
	}
	cronJob.Start()

	return nil
}

func (c *Controller) Shutdown() {
	cronJob.Stop()
}

func (c *Controller) generate() {
	for {
		name, err := c.fileSelector.Select()
		if err != nil {
			c.log.Error("Could not fetch file name: %s", err)
			return
		}
		c.log.Info("Updating sitemap: %s", name)
		// there is not any waiting sitemap updates
		if name == "" {
			return
		}

		c.fileName = name

		els, err := c.fetchElements()
		if err != nil {
			c.log.Error("Could not fetch updated elements: %s", err)
			return
		}

		container := c.buildContainer(els)

		s, err := c.getCurrentSet()
		if err != nil {
			c.log.Error("Could not get current set: %s", err)
			return
		}

		if err := c.updateFile(container, s); err != nil {
			c.log.Error("Could not update file: %s", err)
		}
	}
}

func (c *Controller) fetchElements() ([]*models.SitemapItem, error) {
	key := common.PrepareFileCacheKey(c.fileName)
	redisConn := helper.MustGetRedisConn()
	els := make([]*models.SitemapItem, 0)

	for {
		item, err := redisConn.PopSetMember(key)
		if err != nil && err != redis.ErrNil {
			return els, err
		}

		if item == "" {
			return els, nil
		}

		i := &models.SitemapItem{}

		if err := i.Populate(item); err != nil {
			c.log.Error("Could not update item %s: %s", item, err)
			continue
		}

		els = append(els, i)
	}
}

func (c *Controller) getCurrentSet() (*models.ItemSet, error) {
	// check if this is a new sitemap file or not
	n := fmt.Sprintf("%s.xml", c.fileName)
	if _, err := os.Stat(n); os.IsNotExist(err) {
		return models.NewItemSet(), nil
	}
	input, err := ioutil.ReadFile(n)
	if err != nil {
		return nil, err
	}

	s := models.NewItemSet()
	if err := xml.Unmarshal(input, s); err != nil {
		return nil, err
	}

	return s, nil
}

func (c *Controller) buildContainer(items []*models.SitemapItem) *models.ItemContainer {
	container := models.NewItemContainer()
	for _, v := range items {
		item := v.Definition(config.Get().Uri)
		switch v.Status {
		case models.STATUS_ADD:
			container.Add = append(container.Add, item)
		case models.STATUS_DELETE:
			container.Delete = append(container.Delete, item)
		case models.STATUS_UPDATE:
			container.Update = append(container.Update, item)
		}
	}

	return container
}

func (c *Controller) updateFile(container *models.ItemContainer, set *models.ItemSet) error {
	set.Populate(container)

	return common.XML(set, c.fileName)
}
<commit_msg>Sitemap: Init Sitemap generation cron<commit_after>package generator

import (
	"encoding/xml"
	"fmt"
	"io/ioutil"
	"os"
	"socialapi/config"
	"socialapi/workers/helper"
	"socialapi/workers/sitemap/common"
	"socialapi/workers/sitemap/models"

	"github.com/garyburd/redigo/redis"
	"github.com/koding/logging"
	"github.com/robfig/cron"
)

type Controller struct {
	log          logging.Logger
	fileSelector FileSelector
	fileName     string
}

const (
	SCHEDULE = "0 0-59/30 * * * *"
)

var (
	cronJob *cron.Cron
)

func New(log logging.Logger) (*Controller, error) {
	c := &Controller{
		log:          log,
		fileSelector: CachedFileSelector{},
	}

	return c, c.initCron()
}

func (c *Controller) initCron() error {
	cronJob := cron.New()
	if err := cronJob.AddFunc(SCHEDULE, c.generate); err != nil {
		return err
	}
	cronJob.Start()

	return nil
}

func (c *Controller) Shutdown() {
	cronJob.Stop()
}

func (c *Controller) generate() {
	for {
		name, err := c.fileSelector.Select()
		if err != nil {
			c.log.Error("Could not fetch file name: %s", err)
			return
		}
		c.log.Info("Updating sitemap: %s", name)
		// there is not any waiting sitemap updates
		if name == "" {
			return
		}

		c.fileName = name

		els, err := c.fetchElements()
		if err != nil {
			c.log.Error("Could not fetch updated elements: %s", err)
			return
		}

		container := c.buildContainer(els)

		s, err := c.getCurrentSet()
		if err != nil {
			c.log.Error("Could not get current set: %s", err)
			return
		}

		if err := c.updateFile(container, s); err != nil {
			c.log.Error("Could not update file: %s", err)
		}
	}
}

func (c *Controller) fetchElements() ([]*models.SitemapItem, error) {
	key := common.PrepareFileCacheKey(c.fileName)
	redisConn := helper.MustGetRedisConn()
	els := make([]*models.SitemapItem, 0)

	for {
		item, err := redisConn.PopSetMember(key)
		if err != nil && err != redis.ErrNil {
			return els, err
		}

		if item == "" {
			return els, nil
		}

		i := &models.SitemapItem{}

		if err := i.Populate(item); err != nil {
			c.log.Error("Could not update item %s: %s", item, err)
			continue
		}

		els = append(els, i)
	}
}

func (c *Controller) getCurrentSet() (*models.ItemSet, error) {
	// check if this is a new sitemap file or not
	n := fmt.Sprintf("%s.xml", c.fileName)
	if _, err := os.Stat(n); os.IsNotExist(err) {
		return models.NewItemSet(), nil
	}
	input, err := ioutil.ReadFile(n)
	if err != nil {
		return nil, err
	}

	s := models.NewItemSet()
	if err := xml.Unmarshal(input, s); err != nil {
		return nil, err
	}

	return s, nil
}

func (c *Controller) buildContainer(items []*models.SitemapItem) *models.ItemContainer {
	container := models.NewItemContainer()
	for _, v := range items {
		item := v.Definition(config.Get().Uri)
		switch v.Status {
		case models.STATUS_ADD:
			container.Add = append(container.Add, item)
		case models.STATUS_DELETE:
			container.Delete = append(container.Delete, item)
		case models.STATUS_UPDATE:
			container.Update = append(container.Update, item)
		}
	}

	return container
}

func (c *Controller) updateFile(container *models.ItemContainer, set *models.ItemSet) error {
	set.Populate(container)

	return common.XML(set, c.fileName)
}
<|endoftext|>


<commit_before>//
// store.go
//

package srnd

import (
  "github.com/majestrate/srndv2/src/nacl"
  "bufio"
  "bytes"
  "crypto/sha512"
  "errors"
  "io"
  "log"
  "mime"
  "mime/multipart"
  "net/mail"
  "os"
  "os/exec"
  "path/filepath"
)

type ArticleStore interface {
  MessageReader
  MessageWriter

  // get the filepath for an attachment
  AttachmentFilepath(fname string) string
  // get the filepath for an attachment's thumbnail
  ThumbnailFilepath(fname string) string
  // do we have this article?
  HasArticle(msgid string) bool
  // create a file for a message
  CreateFile(msgid string) io.WriteCloser
  // create a file for a temp message, returns nil if it's already open
  CreateTempFile(msgid string) io.WriteCloser
  // get the filename of a message
  GetFilename(msgid string) string
  // get the filename of a temp message
  GetTempFilename(msgid string) string
  // Get a message given its messageid
  GetMessage(msgid string) NNTPMessage
  // get a temp message given its messageid
  // temp message is deleted once read
  ReadTempMessage(msgid string) NNTPMessage
  // store a post
  StorePost(nntp NNTPMessage) error
  // get article headers only
  GetHeaders(msgid string) ArticleHeaders

  // get our temp directory for articles
  TempDir() string
}
type articleStore struct {
  directory string
  temp string
  attachments string
  thumbs string
  database Database
  convert_path string
}

func createArticleStore(config map[string]string, database Database) ArticleStore {
  store := articleStore{
    directory: config["store_dir"],
    temp: config["incoming_dir"],
    attachments: config["attachments_dir"],
    thumbs: config["thumbs_dir"],
    convert_path: config["convert_bin"],
    database: database,
  }
  store.Init()
  return store
}

func (self articleStore) TempDir() string {
  return self.temp
}

// initialize article store
func (self articleStore) Init() {
  EnsureDir(self.directory)
  EnsureDir(self.temp)
  EnsureDir(self.attachments)
  EnsureDir(self.thumbs)
  if ! CheckFile(self.convert_path) {
    log.Fatal("cannot find executable for convert:", self.convert_path, "not found")
  }
}

func (self articleStore) GenerateThumbnail(fname string) error {
  outfname := self.ThumbnailFilepath(fname)
  infname := self.AttachmentFilepath(fname)
  cmd := exec.Command(self.convert_path, "-thumbnail", "200", infname, outfname)
  exec_out, err := cmd.CombinedOutput()
  if err != nil {
    log.Println("error generating thumbnail", string(exec_out))
  }
  return err
}

func (self articleStore) ReadMessage(r io.Reader) (NNTPMessage, error) {
  return read_message(r)
}

func (self articleStore) StorePost(nntp NNTPMessage) (err error) {
  f := self.CreateFile(nntp.MessageID())
  if f != nil {
    err = self.WriteMessage(nntp, f)
    f.Close()
  }

  nntp_inner := nntp.Signed()
  if nntp_inner == nil {
    // no inner article
    // store the data in the article
    self.database.RegisterArticle(nntp)
    for _, att := range nntp.Attachments() {
      // save attachments
      go self.saveAttachment(att)
    }
  } else {
    // record a tripcode
    self.database.RegisterSigned(nntp.MessageID(), nntp.Pubkey())
    // we have inner data
    // store the signed data
    self.database.RegisterArticle(nntp_inner)
    for _, att := range nntp_inner.Attachments() {
      go self.saveAttachment(att)
    }
  }
  return
}

// save an attachment
func (self articleStore) saveAttachment(att NNTPAttachment) {
  var err error
  var f io.WriteCloser
  fpath := att.Filepath()
  upload := self.AttachmentFilepath(fpath)
  thumb := self.ThumbnailFilepath(fpath)
  if CheckFile(upload) {
    log.Println("already have file", fpath)
    if ! CheckFile(thumb) && att.NeedsThumbnail() {
      log.Println("create thumbnail for", fpath)
      err = self.GenerateThumbnail(fpath)
      if err != nil {
        log.Println("failed to generate thumbnail", err)
      }
    }
    return
  }
  // save attachment
  log.Println("save attachment", att.Filename(), "to", upload)
  f, err = os.Create(upload)
  if err == nil {
    err = att.WriteTo(f)
    f.Close()
  }
  if err != nil {
    log.Println("did not save attachment", err)
    return
  }

  // generate thumbanils
  if att.NeedsThumbnail() {
    log.Println("create thumbnail for", fpath)
    err = self.GenerateThumbnail(fpath)
    if err != nil {
      log.Println("failed to generate thumbnail", err)
    }
  }
}

// eh this isn't really needed is it?
func (self articleStore) WriteMessage(nntp NNTPMessage, wr io.Writer) (err error) {
  return nntp.WriteTo(wr, "\n")
}

// get the filepath for an attachment
func (self articleStore) AttachmentFilepath(fname string) string {
  return filepath.Join(self.attachments, fname)
}

// get the filepath for a thumbanil
func (self articleStore) ThumbnailFilepath(fname string) string {
  return filepath.Join(self.thumbs, fname)
}

// create a file for this article
func (self articleStore) CreateFile(messageID string) io.WriteCloser {
  fname := self.GetFilename(messageID)
  file, err := os.Create(fname)
  if err != nil {
    log.Println("cannot open file", fname)
    return nil
  }
  return file
}

// create a temp file for inboud articles
func (self articleStore) CreateTempFile(messageID string) io.WriteCloser {
  fname := self.GetTempFilename(messageID)
  if CheckFile(fname) {
    log.Println(fname, "already open")
    return nil
  }
  file, err := os.Create(fname)
  if err != nil {
    log.Println("cannot open file", fname)
    return nil
  }
  return file
}

// return true if we have an article
func (self articleStore) HasArticle(messageID string) bool {
  return self.database.HasArticleLocal(messageID)
}

// get the filename for this article
func (self articleStore) GetFilename(messageID string) string {
  if ! ValidMessageID(messageID) {
    log.Println("!!! bug: tried to open invalid message", messageID, "!!!")
    return ""
  }
  return filepath.Join(self.directory, messageID)
}

// get the filename for this article
func (self articleStore) GetTempFilename(messageID string) string {
  if ! ValidMessageID(messageID) {
    log.Println("!!! bug: tried to open invalid temp message", messageID, "!!!")
    return ""
  }
  return filepath.Join(self.temp, messageID)
}

// loads temp message and deletes old article
func (self articleStore) ReadTempMessage(messageID string) NNTPMessage {
  fname := self.GetTempFilename(messageID)
  nntp := self.readfile(fname)
  DelFile(fname)
  return nntp
}

// read a file give filepath
func (self articleStore) readfile(fname string) NNTPMessage {

  file, err := os.Open(fname)
  if err != nil {
    log.Println("store cannot open file",fname)
    return nil
  }
  message, err := self.ReadMessage(file)
  file.Close()
  if err == nil {
    return message
  }

  log.Println("failed to load file", fname)
  return nil
}

// get the replies for a thread
func (self articleStore) GetThreadReplies(messageID string, last int) []NNTPMessage {
  var repls []NNTPMessage
  if self.database.ThreadHasReplies(messageID) {
    rpls := self.database.GetThreadReplies(messageID, last)
    if rpls == nil {
      return repls
    }
    for _, rpl := range rpls {
      msg := self.GetMessage(rpl)
      if msg == nil {
        log.Println("cannot get message", rpl)
      } else {
        repls = append(repls, msg)
      }
    }
  }
  return repls
}

// load an article
// return nil on failure
func (self articleStore) GetMessage(messageID string) NNTPMessage {
  return self.readfile(self.GetFilename(messageID))
}

// get article with headers only
func (self articleStore) GetHeaders(messageID string) ArticleHeaders {
  // TODO: don't load the entire body
  nntp := self.readfile(self.GetFilename(messageID))
  if nntp == nil {
    return nil
  }
  return nntp.Headers()
}

func read_message(r io.Reader) (NNTPMessage, error) {

  msg, err := mail.ReadMessage(r)
  var nntp nntpArticle

  if err == nil {
    nntp.headers = ArticleHeaders(msg.Header)
    content_type := nntp.ContentType()
    media_type, params, err := mime.ParseMediaType(content_type)
    if err != nil {
      log.Println("failed to parse media type", err, "for mime", content_type)
      return nil, err
    }
    boundary, ok := params["boundary"]
    if ok {
      partReader := multipart.NewReader(msg.Body, boundary)
      for {
        part, err := partReader.NextPart()
        if err == io.EOF {
          return nntp, nil
        } else if err == nil {
          hdr := part.Header
          // get content type of part
          part_type := hdr.Get("Content-Type")
          // parse content type
          media_type, _, err = mime.ParseMediaType(part_type)
          if err == nil {
            if media_type == "text/plain" {
              att := readAttachmentFromMimePart(part)
              if att != nil {
                nntp.message = att.(nntpAttachment)
                nntp.message.header.Set("Content-Type", part_type)
              }
            } else {
              // non plaintext gets added to attachments
              att := readAttachmentFromMimePart(part)
              if att != nil {
                nntp = nntp.Attach(att).(nntpArticle)
              }
            }
          } else {
            log.Println("part has no content type", err)
          }
          part.Close()
        } else {
          log.Println("failed to load part! ", err)
          return nil, err
        }
      }

    } else if media_type == "message/rfc822" {
      // tripcoded message
      sig := nntp.headers.Get("X-Signature-Ed25519-Sha512", "")
      pk := nntp.Pubkey()
      if pk == "" || sig == "" {
        log.Println("invalid sig or pubkey", sig, pk)
        return nil, errors.New("invalid headers")
      }
      log.Printf("got signed message from %s", pk)
      pk_bytes := unhex(pk)
      sig_bytes := unhex(sig)
      r := bufio.NewReader(msg.Body)
      crlf := []byte{13,10}
      for {
        line, err := r.ReadBytes('\n')
        if err == io.EOF {
          break
        }
        nntp.signedPart.body.Write(line[:len(line)-1])
        nntp.signedPart.body.Write(crlf)
      }
      if nntp.signedPart.body.Len() < 2 {
        log.Println("signed body is too small")
      } else {
        body := nntp.signedPart.body.Bytes()[:nntp.signedPart.body.Len()-2]
        body_hash := sha512.Sum512(body)
        log.Printf("hash=%s", hexify(body_hash[:]))
        if nacl.CryptoVerifyFucky(body_hash[:], sig_bytes, pk_bytes) {
          log.Println("signature is valid :^)")
          return nntp, nil
        } else {
          log.Println("!!!signature is invalid!!!")
        }
      }
    } else {
      // plaintext attachment
      var buff bytes.Buffer
      _, err = io.Copy(&buff, msg.Body)
      nntp.message = createPlaintextAttachment(buff.String())
      return nntp, err
    }
  } else {
    log.Println("failed to read message", err)
    return nil, err
  }
  return nntp, err
}

<commit_msg>have store check filesystem instead of database<commit_after>//
// store.go
//

package srnd

import (
  "github.com/majestrate/srndv2/src/nacl"
  "bufio"
  "bytes"
  "crypto/sha512"
  "errors"
  "io"
  "log"
  "mime"
  "mime/multipart"
  "net/mail"
  "os"
  "os/exec"
  "path/filepath"
)

type ArticleStore interface {
  MessageReader
  MessageWriter

  // get the filepath for an attachment
  AttachmentFilepath(fname string) string
  // get the filepath for an attachment's thumbnail
  ThumbnailFilepath(fname string) string
  // do we have this article?
  HasArticle(msgid string) bool
  // create a file for a message
  CreateFile(msgid string) io.WriteCloser
  // create a file for a temp message, returns nil if it's already open
  CreateTempFile(msgid string) io.WriteCloser
  // get the filename of a message
  GetFilename(msgid string) string
  // get the filename of a temp message
  GetTempFilename(msgid string) string
  // Get a message given its messageid
  GetMessage(msgid string) NNTPMessage
  // get a temp message given its messageid
  // temp message is deleted once read
  ReadTempMessage(msgid string) NNTPMessage
  // store a post
  StorePost(nntp NNTPMessage) error
  // get article headers only
  GetHeaders(msgid string) ArticleHeaders

  // get our temp directory for articles
  TempDir() string
}
type articleStore struct {
  directory string
  temp string
  attachments string
  thumbs string
  database Database
  convert_path string
}

func createArticleStore(config map[string]string, database Database) ArticleStore {
  store := articleStore{
    directory: config["store_dir"],
    temp: config["incoming_dir"],
    attachments: config["attachments_dir"],
    thumbs: config["thumbs_dir"],
    convert_path: config["convert_bin"],
    database: database,
  }
  store.Init()
  return store
}

func (self articleStore) TempDir() string {
  return self.temp
}

// initialize article store
func (self articleStore) Init() {
  EnsureDir(self.directory)
  EnsureDir(self.temp)
  EnsureDir(self.attachments)
  EnsureDir(self.thumbs)
  if ! CheckFile(self.convert_path) {
    log.Fatal("cannot find executable for convert:", self.convert_path, "not found")
  }
}

func (self articleStore) GenerateThumbnail(fname string) error {
  outfname := self.ThumbnailFilepath(fname)
  infname := self.AttachmentFilepath(fname)
  cmd := exec.Command(self.convert_path, "-thumbnail", "200", infname, outfname)
  exec_out, err := cmd.CombinedOutput()
  if err != nil {
    log.Println("error generating thumbnail", string(exec_out))
  }
  return err
}

func (self articleStore) ReadMessage(r io.Reader) (NNTPMessage, error) {
  return read_message(r)
}

func (self articleStore) StorePost(nntp NNTPMessage) (err error) {
  f := self.CreateFile(nntp.MessageID())
  if f != nil {
    err = self.WriteMessage(nntp, f)
    f.Close()
  }

  nntp_inner := nntp.Signed()
  if nntp_inner == nil {
    // no inner article
    // store the data in the article
    self.database.RegisterArticle(nntp)
    for _, att := range nntp.Attachments() {
      // save attachments
      go self.saveAttachment(att)
    }
  } else {
    // record a tripcode
    self.database.RegisterSigned(nntp.MessageID(), nntp.Pubkey())
    // we have inner data
    // store the signed data
    self.database.RegisterArticle(nntp_inner)
    for _, att := range nntp_inner.Attachments() {
      go self.saveAttachment(att)
    }
  }
  return
}

// save an attachment
func (self articleStore) saveAttachment(att NNTPAttachment) {
  var err error
  var f io.WriteCloser
  fpath := att.Filepath()
  upload := self.AttachmentFilepath(fpath)
  thumb := self.ThumbnailFilepath(fpath)
  if CheckFile(upload) {
    log.Println("already have file", fpath)
    if ! CheckFile(thumb) && att.NeedsThumbnail() {
      log.Println("create thumbnail for", fpath)
      err = self.GenerateThumbnail(fpath)
      if err != nil {
        log.Println("failed to generate thumbnail", err)
      }
    }
    return
  }
  // save attachment
  log.Println("save attachment", att.Filename(), "to", upload)
  f, err = os.Create(upload)
  if err == nil {
    err = att.WriteTo(f)
    f.Close()
  }
  if err != nil {
    log.Println("did not save attachment", err)
    return
  }

  // generate thumbanils
  if att.NeedsThumbnail() {
    log.Println("create thumbnail for", fpath)
    err = self.GenerateThumbnail(fpath)
    if err != nil {
      log.Println("failed to generate thumbnail", err)
    }
  }
}

// eh this isn't really needed is it?
func (self articleStore) WriteMessage(nntp NNTPMessage, wr io.Writer) (err error) {
  return nntp.WriteTo(wr, "\n")
}

// get the filepath for an attachment
func (self articleStore) AttachmentFilepath(fname string) string {
  return filepath.Join(self.attachments, fname)
}

// get the filepath for a thumbanil
func (self articleStore) ThumbnailFilepath(fname string) string {
  return filepath.Join(self.thumbs, fname)
}

// create a file for this article
func (self articleStore) CreateFile(messageID string) io.WriteCloser {
  fname := self.GetFilename(messageID)
  file, err := os.Create(fname)
  if err != nil {
    log.Println("cannot open file", fname)
    return nil
  }
  return file
}

// create a temp file for inboud articles
func (self articleStore) CreateTempFile(messageID string) io.WriteCloser {
  fname := self.GetTempFilename(messageID)
  if CheckFile(fname) {
    log.Println(fname, "already open")
    return nil
  }
  file, err := os.Create(fname)
  if err != nil {
    log.Println("cannot open file", fname)
    return nil
  }
  return file
}

// return true if we have an article
func (self articleStore) HasArticle(messageID string) bool {
  return CheckFile(self.GetFilename(messageID))
}

// get the filename for this article
func (self articleStore) GetFilename(messageID string) string {
  if ! ValidMessageID(messageID) {
    log.Println("!!! bug: tried to open invalid message", messageID, "!!!")
    return ""
  }
  return filepath.Join(self.directory, messageID)
}

// get the filename for this article
func (self articleStore) GetTempFilename(messageID string) string {
  if ! ValidMessageID(messageID) {
    log.Println("!!! bug: tried to open invalid temp message", messageID, "!!!")
    return ""
  }
  return filepath.Join(self.temp, messageID)
}

// loads temp message and deletes old article
func (self articleStore) ReadTempMessage(messageID string) NNTPMessage {
  fname := self.GetTempFilename(messageID)
  nntp := self.readfile(fname)
  DelFile(fname)
  return nntp
}

// read a file give filepath
func (self articleStore) readfile(fname string) NNTPMessage {

  file, err := os.Open(fname)
  if err != nil {
    log.Println("store cannot open file",fname)
    return nil
  }
  message, err := self.ReadMessage(file)
  file.Close()
  if err == nil {
    return message
  }

  log.Println("failed to load file", fname)
  return nil
}

// get the replies for a thread
func (self articleStore) GetThreadReplies(messageID string, last int) []NNTPMessage {
  var repls []NNTPMessage
  if self.database.ThreadHasReplies(messageID) {
    rpls := self.database.GetThreadReplies(messageID, last)
    if rpls == nil {
      return repls
    }
    for _, rpl := range rpls {
      msg := self.GetMessage(rpl)
      if msg == nil {
        log.Println("cannot get message", rpl)
      } else {
        repls = append(repls, msg)
      }
    }
  }
  return repls
}

// load an article
// return nil on failure
func (self articleStore) GetMessage(messageID string) NNTPMessage {
  return self.readfile(self.GetFilename(messageID))
}

// get article with headers only
func (self articleStore) GetHeaders(messageID string) ArticleHeaders {
  // TODO: don't load the entire body
  nntp := self.readfile(self.GetFilename(messageID))
  if nntp == nil {
    return nil
  }
  return nntp.Headers()
}

func read_message(r io.Reader) (NNTPMessage, error) {

  msg, err := mail.ReadMessage(r)
  var nntp nntpArticle

  if err == nil {
    nntp.headers = ArticleHeaders(msg.Header)
    content_type := nntp.ContentType()
    media_type, params, err := mime.ParseMediaType(content_type)
    if err != nil {
      log.Println("failed to parse media type", err, "for mime", content_type)
      return nil, err
    }
    boundary, ok := params["boundary"]
    if ok {
      partReader := multipart.NewReader(msg.Body, boundary)
      for {
        part, err := partReader.NextPart()
        if err == io.EOF {
          return nntp, nil
        } else if err == nil {
          hdr := part.Header
          // get content type of part
          part_type := hdr.Get("Content-Type")
          // parse content type
          media_type, _, err = mime.ParseMediaType(part_type)
          if err == nil {
            if media_type == "text/plain" {
              att := readAttachmentFromMimePart(part)
              if att != nil {
                nntp.message = att.(nntpAttachment)
                nntp.message.header.Set("Content-Type", part_type)
              }
            } else {
              // non plaintext gets added to attachments
              att := readAttachmentFromMimePart(part)
              if att != nil {
                nntp = nntp.Attach(att).(nntpArticle)
              }
            }
          } else {
            log.Println("part has no content type", err)
          }
          part.Close()
        } else {
          log.Println("failed to load part! ", err)
          return nil, err
        }
      }

    } else if media_type == "message/rfc822" {
      // tripcoded message
      sig := nntp.headers.Get("X-Signature-Ed25519-Sha512", "")
      pk := nntp.Pubkey()
      if pk == "" || sig == "" {
        log.Println("invalid sig or pubkey", sig, pk)
        return nil, errors.New("invalid headers")
      }
      log.Printf("got signed message from %s", pk)
      pk_bytes := unhex(pk)
      sig_bytes := unhex(sig)
      r := bufio.NewReader(msg.Body)
      crlf := []byte{13,10}
      for {
        line, err := r.ReadBytes('\n')
        if err == io.EOF {
          break
        }
        nntp.signedPart.body.Write(line[:len(line)-1])
        nntp.signedPart.body.Write(crlf)
      }
      if nntp.signedPart.body.Len() < 2 {
        log.Println("signed body is too small")
      } else {
        body := nntp.signedPart.body.Bytes()[:nntp.signedPart.body.Len()-2]
        body_hash := sha512.Sum512(body)
        log.Printf("hash=%s", hexify(body_hash[:]))
        if nacl.CryptoVerifyFucky(body_hash[:], sig_bytes, pk_bytes) {
          log.Println("signature is valid :^)")
          return nntp, nil
        } else {
          log.Println("!!!signature is invalid!!!")
        }
      }
    } else {
      // plaintext attachment
      var buff bytes.Buffer
      _, err = io.Copy(&buff, msg.Body)
      nntp.message = createPlaintextAttachment(buff.String())
      return nntp, err
    }
  } else {
    log.Println("failed to read message", err)
    return nil, err
  }
  return nntp, err
}

<|endoftext|>