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	"fmt"
20	"io"
21	"sort"
22	"strings"
23)
24
25func peSymbolsFromFile(r io.ReaderAt) (*File, error) {
26	peFile, err := pe.NewFile(r)
27	if err != nil {
28		return nil, cantParseError{err}
29	}
30
31	return extractPESymbols(peFile)
32}
33
34func extractPESymbols(peFile *pe.File) (*File, error) {
35	var prefix string
36	if peFile.FileHeader.Machine == pe.IMAGE_FILE_MACHINE_I386 {
37		// symbols in win32 exes seem to be prefixed with an underscore
38		prefix = "_"
39	}
40
41	symbols := peFile.Symbols
42	sort.SliceStable(symbols, func(i, j int) bool {
43		if symbols[i].SectionNumber != symbols[j].SectionNumber {
44			return symbols[i].SectionNumber < symbols[j].SectionNumber
45		}
46		return symbols[i].Value < symbols[j].Value
47	})
48
49	file := &File{}
50
51	for _, section := range peFile.Sections {
52		file.Sections = append(file.Sections, &Section{
53			Name:   section.Name,
54			Addr:   uint64(section.VirtualAddress),
55			Offset: uint64(section.Offset),
56			Size:   uint64(section.VirtualSize),
57		})
58	}
59
60	for _, symbol := range symbols {
61		if symbol.SectionNumber > 0 {
62			file.Symbols = append(file.Symbols, &Symbol{
63				Name: strings.TrimPrefix(symbol.Name, prefix),
64				// PE symbol value is the offset of the symbol into the section
65				Addr: uint64(symbol.Value),
66				// PE symbols don't have size information
67				Size:    0,
68				Section: file.Sections[symbol.SectionNumber-1],
69			})
70		}
71	}
72
73	return file, nil
74}
75
76func dumpPESymbols(r io.ReaderAt) error {
77	peFile, err := pe.NewFile(r)
78	if err != nil {
79		return cantParseError{err}
80	}
81
82	fmt.Println("&pe.File{")
83	fmt.Println("\tFileHeader: pe.FileHeader{")
84	fmt.Printf("\t\tMachine: %#v,\n", peFile.FileHeader.Machine)
85	fmt.Println("\t},")
86
87	fmt.Println("\tSections: []*pe.Section{")
88	for _, section := range peFile.Sections {
89		fmt.Printf("\t\t&pe.Section{SectionHeader: %#v},\n", section.SectionHeader)
90	}
91	fmt.Println("\t},")
92
93	fmt.Println("\tSymbols: []*pe.Symbol{")
94	for _, symbol := range peFile.Symbols {
95		fmt.Printf("\t\t%#v,\n", symbol)
96	}
97	fmt.Println("\t},")
98
99	fmt.Println("}")
100
101	return nil
102}
103