1// Copyright 2018 Google Inc. All rights reserved.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15package symbol_inject
16
17import (
18	"debug/pe"
19	"strconv"
20	"testing"
21)
22
23func TestPESymbolTable(t *testing.T) {
24	testCases := []struct {
25		file         *pe.File
26		symbol       string
27		offset, size uint64
28	}{
29		{
30			file:   peSymbolTable1,
31			symbol: "soong_build_number",
32			offset: 0x2420,
33			size:   128,
34		},
35		{
36			file:   peSymbolTable2,
37			symbol: "symbol1",
38			offset: 0x2420,
39			size:   128,
40		},
41		{
42			file:   peSymbolTable2,
43			symbol: "symbol2",
44			offset: 0x24a0,
45			size:   128,
46		},
47		{
48			// Test when symbol has the same value as the target symbol but is located afterwards in the list
49			file: &pe.File{
50				FileHeader: pe.FileHeader{
51					Machine: pe.IMAGE_FILE_MACHINE_I386,
52				},
53				Sections: []*pe.Section{
54					&pe.Section{SectionHeader: pe.SectionHeader{Name: ".text", VirtualSize: 0x15e83c, VirtualAddress: 0x1000, Size: 0x15ea00, Offset: 0x600, PointerToRelocations: 0x0, PointerToLineNumbers: 0x0, NumberOfRelocations: 0x0, NumberOfLineNumbers: 0x0, Characteristics: 0x60500060}},
55					&pe.Section{SectionHeader: pe.SectionHeader{Name: ".data", VirtualSize: 0x6a58, VirtualAddress: 0x160000, Size: 0x6c00, Offset: 0x15f000, PointerToRelocations: 0x0, PointerToLineNumbers: 0x0, NumberOfRelocations: 0x0, NumberOfLineNumbers: 0x0, Characteristics: 0xc0600040}},
56				},
57				Symbols: []*pe.Symbol{
58					&pe.Symbol{Name: "_soong_build_number", Value: 0x20, SectionNumber: 2, Type: 0x0, StorageClass: 0x2},
59					&pe.Symbol{Name: ".data", Value: 0x20, SectionNumber: 2, Type: 0x0, StorageClass: 0x3},
60					&pe.Symbol{Name: "_adb_device_banner", Value: 0xa0, SectionNumber: 2, Type: 0x0, StorageClass: 0x2},
61				},
62			},
63			symbol: "soong_build_number",
64			offset: 0x15f020,
65			size:   128,
66		},
67		{
68			// Test when symbol has nothing after it
69			file: &pe.File{
70				FileHeader: pe.FileHeader{
71					Machine: pe.IMAGE_FILE_MACHINE_AMD64,
72				},
73				Sections: []*pe.Section{
74					&pe.Section{SectionHeader: pe.SectionHeader{Name: ".text", VirtualSize: 0x1cc0, VirtualAddress: 0x1000, Size: 0x1e00, Offset: 0x600, PointerToRelocations: 0x0, PointerToLineNumbers: 0x0, NumberOfRelocations: 0x0, NumberOfLineNumbers: 0x0, Characteristics: 0x60500020}},
75					&pe.Section{SectionHeader: pe.SectionHeader{Name: ".data", VirtualSize: 0xa0, VirtualAddress: 0x3000, Size: 0x200, Offset: 0x2400, PointerToRelocations: 0x0, PointerToLineNumbers: 0x0, NumberOfRelocations: 0x0, NumberOfLineNumbers: 0x0, Characteristics: 0xc0600040}},
76				},
77				Symbols: []*pe.Symbol{
78					&pe.Symbol{Name: "soong_build_number", Value: 0x20, SectionNumber: 2, Type: 0x0, StorageClass: 0x2},
79				},
80			},
81			symbol: "soong_build_number",
82			offset: 0x2420,
83			size:   128,
84		},
85		{
86			// Test when symbol has a symbol in a different section after it
87			file: &pe.File{
88				FileHeader: pe.FileHeader{
89					Machine: pe.IMAGE_FILE_MACHINE_AMD64,
90				},
91				Sections: []*pe.Section{
92					&pe.Section{SectionHeader: pe.SectionHeader{Name: ".text", VirtualSize: 0x1cc0, VirtualAddress: 0x1000, Size: 0x1e00, Offset: 0x600, PointerToRelocations: 0x0, PointerToLineNumbers: 0x0, NumberOfRelocations: 0x0, NumberOfLineNumbers: 0x0, Characteristics: 0x60500020}},
93					&pe.Section{SectionHeader: pe.SectionHeader{Name: ".data", VirtualSize: 0xa0, VirtualAddress: 0x3000, Size: 0x200, Offset: 0x2400, PointerToRelocations: 0x0, PointerToLineNumbers: 0x0, NumberOfRelocations: 0x0, NumberOfLineNumbers: 0x0, Characteristics: 0xc0600040}},
94					&pe.Section{SectionHeader: pe.SectionHeader{Name: ".rdata", VirtualSize: 0x5e0, VirtualAddress: 0x4000, Size: 0x600, Offset: 0x2600, PointerToRelocations: 0x0, PointerToLineNumbers: 0x0, NumberOfRelocations: 0x0, NumberOfLineNumbers: 0x0, Characteristics: 0x40500040}},
95				},
96				Symbols: []*pe.Symbol{
97					&pe.Symbol{Name: "soong_build_number", Value: 0x20, SectionNumber: 2, Type: 0x0, StorageClass: 0x2},
98					&pe.Symbol{Name: "_adb_device_banner", Value: 0x30, SectionNumber: 3, Type: 0x0, StorageClass: 0x2},
99				},
100			},
101			symbol: "soong_build_number",
102			offset: 0x2420,
103			size:   128,
104		},
105		{
106			// Test when symbols are out of order
107			file: &pe.File{
108				FileHeader: pe.FileHeader{
109					Machine: pe.IMAGE_FILE_MACHINE_AMD64,
110				},
111				Sections: []*pe.Section{
112					&pe.Section{SectionHeader: pe.SectionHeader{Name: ".text", VirtualSize: 0x1cc0, VirtualAddress: 0x1000, Size: 0x1e00, Offset: 0x600, PointerToRelocations: 0x0, PointerToLineNumbers: 0x0, NumberOfRelocations: 0x0, NumberOfLineNumbers: 0x0, Characteristics: 0x60500020}},
113					&pe.Section{SectionHeader: pe.SectionHeader{Name: ".data", VirtualSize: 0x120, VirtualAddress: 0x3000, Size: 0x200, Offset: 0x2400, PointerToRelocations: 0x0, PointerToLineNumbers: 0x0, NumberOfRelocations: 0x0, NumberOfLineNumbers: 0x0, Characteristics: 0xc0600040}},
114				},
115				Symbols: []*pe.Symbol{
116					&pe.Symbol{Name: "_adb_device_banner", Value: 0xa0, SectionNumber: 2, Type: 0x0, StorageClass: 0x2},
117					&pe.Symbol{Name: "soong_build_number", Value: 0x20, SectionNumber: 2, Type: 0x0, StorageClass: 0x2},
118				},
119			},
120			symbol: "soong_build_number",
121			offset: 0x2420,
122			size:   128,
123		},
124	}
125
126	for i, testCase := range testCases {
127		t.Run(strconv.Itoa(i), func(t *testing.T) {
128			file, err := extractPESymbols(testCase.file)
129			if err != nil {
130				t.Error(err.Error())
131			}
132			offset, size, err := findSymbol(file, testCase.symbol)
133			if err != nil {
134				t.Error(err.Error())
135			}
136			if offset != testCase.offset || size != testCase.size {
137				t.Errorf("expected %x:%x, got %x:%x", testCase.offset, testCase.size, offset, size)
138			}
139		})
140	}
141}
142