| Colin Cross | 5498f85 | 2018-01-03 23:39:54 -0800 | [diff] [blame] | 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 |  | 
| Dan Willemsen | 2249dc8 | 2018-10-15 00:35:59 -0700 | [diff] [blame] | 15 | package symbol_inject | 
| Colin Cross | 5498f85 | 2018-01-03 23:39:54 -0800 | [diff] [blame] | 16 |  | 
|  | 17 | import ( | 
|  | 18 | "debug/pe" | 
|  | 19 | "fmt" | 
|  | 20 | "io" | 
|  | 21 | "sort" | 
| Colin Cross | 8673b5b | 2018-03-01 11:20:25 -0800 | [diff] [blame] | 22 | "strings" | 
| Colin Cross | 5498f85 | 2018-01-03 23:39:54 -0800 | [diff] [blame] | 23 | ) | 
|  | 24 |  | 
| Colin Cross | 8673b5b | 2018-03-01 11:20:25 -0800 | [diff] [blame] | 25 | func peSymbolsFromFile(r io.ReaderAt) (*File, error) { | 
| Colin Cross | 5498f85 | 2018-01-03 23:39:54 -0800 | [diff] [blame] | 26 | peFile, err := pe.NewFile(r) | 
|  | 27 | if err != nil { | 
| Colin Cross | 8673b5b | 2018-03-01 11:20:25 -0800 | [diff] [blame] | 28 | return nil, cantParseError{err} | 
| Colin Cross | 5498f85 | 2018-01-03 23:39:54 -0800 | [diff] [blame] | 29 | } | 
|  | 30 |  | 
| Colin Cross | 8673b5b | 2018-03-01 11:20:25 -0800 | [diff] [blame] | 31 | return extractPESymbols(peFile) | 
|  | 32 | } | 
|  | 33 |  | 
|  | 34 | func extractPESymbols(peFile *pe.File) (*File, error) { | 
|  | 35 | var prefix string | 
| Colin Cross | 64c6d4b | 2018-02-23 17:10:34 -0800 | [diff] [blame] | 36 | if peFile.FileHeader.Machine == pe.IMAGE_FILE_MACHINE_I386 { | 
|  | 37 | // symbols in win32 exes seem to be prefixed with an underscore | 
| Colin Cross | 8673b5b | 2018-03-01 11:20:25 -0800 | [diff] [blame] | 38 | prefix = "_" | 
| Colin Cross | 64c6d4b | 2018-02-23 17:10:34 -0800 | [diff] [blame] | 39 | } | 
|  | 40 |  | 
| Colin Cross | dfce764 | 2018-02-28 13:05:39 -0800 | [diff] [blame] | 41 | symbols := peFile.Symbols | 
| Colin Cross | 8673b5b | 2018-03-01 11:20:25 -0800 | [diff] [blame] | 42 | sort.SliceStable(symbols, func(i, j int) bool { | 
| Colin Cross | dfce764 | 2018-02-28 13:05:39 -0800 | [diff] [blame] | 43 | if symbols[i].SectionNumber != symbols[j].SectionNumber { | 
|  | 44 | return symbols[i].SectionNumber < symbols[j].SectionNumber | 
| Colin Cross | 5498f85 | 2018-01-03 23:39:54 -0800 | [diff] [blame] | 45 | } | 
| Colin Cross | dfce764 | 2018-02-28 13:05:39 -0800 | [diff] [blame] | 46 | return symbols[i].Value < symbols[j].Value | 
| Colin Cross | 5498f85 | 2018-01-03 23:39:54 -0800 | [diff] [blame] | 47 | }) | 
|  | 48 |  | 
| Colin Cross | 8673b5b | 2018-03-01 11:20:25 -0800 | [diff] [blame] | 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 |  | 
| Colin Cross | dfce764 | 2018-02-28 13:05:39 -0800 | [diff] [blame] | 60 | for _, symbol := range symbols { | 
| Colin Cross | 8673b5b | 2018-03-01 11:20:25 -0800 | [diff] [blame] | 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], | 
| Colin Cross | dfce764 | 2018-02-28 13:05:39 -0800 | [diff] [blame] | 69 | }) | 
| Colin Cross | 5498f85 | 2018-01-03 23:39:54 -0800 | [diff] [blame] | 70 | } | 
|  | 71 | } | 
|  | 72 |  | 
| Colin Cross | 8673b5b | 2018-03-01 11:20:25 -0800 | [diff] [blame] | 73 | return file, nil | 
|  | 74 | } | 
|  | 75 |  | 
|  | 76 | func 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 | 
| Colin Cross | 5498f85 | 2018-01-03 23:39:54 -0800 | [diff] [blame] | 102 | } |